feat: add running operations tracking to prevent unsafe termination (#10)

## Running Operations Tracking

### **What We Built:**
1. hasRunningOperations() method - Clean API to check if critical
operations are running
2. Internal tracking system - addRunningOperation() /
removeRunningOperation() with Set-based storage
3. Automatic cleanup - try/finally blocks ensure operations are removed
even on errors

### **Where We Added Tracking (Only 2 Operations):**

✅ ctx.step:
• **Executes user code** 
• **Critical for workflow state** - if interrupted, workflow state
becomes inconsistent

✅ ctx.waitForCondition:
• **Executes user condition checks** 
• **Critical for workflow logic** - if interrupted, condition evaluation
becomes unreliable

### **Why NOT Other Operations:**

❌ ctx.runInChildContext - Safe to interrupt, just context management
❌ ctx.wait - Simple timer, no user code execution
❌ ctx.createCallback - Just creates callback config, no execution
❌ ctx.waitForCallback - Uses ctx.step internally, so already tracked
automatically
❌ ctx.map/parallel/executeConcurrently - Use ctx.runInChildContext
internally, which is safe to interrupt

### **Key Insight:**
We only needed to add tracking to the 2 fundamental operations that
execute user code:
• ctx.step - Direct user code execution
• ctx.waitForCondition - User condition evaluation

All other operations either:
• Are safe to interrupt (context/config operations)
• Delegate to tracked operations (so get protection automatically)

### **The Result:**
```typescript
// Safe termination check
if (durableContext.hasRunningOperations()) {
  return; // Critical user code is running - wait
}
terminate(); // Safe to terminate
```

Surgical precision - Only the 2 operations that actually execute user
code are tracked! 🎯

Co-authored-by: Pooya Paridel <parpooya@amazon.com>

Author: ParidelPooya

packages/lambda-durable-functions-sdk-js/src/context/durable-context/durable-context.test.ts

@@ -138,10 +138,21 @@ describe("Durable Context", () => {
       mockParentContext,
       expect.any(Function), // createStepId
       expect.any(Function), // createContextLogger
+      expect.any(Function), // addRunningOperation
+      expect.any(Function), // removeRunningOperation
     );
     expect(mockStepHandler).toHaveBeenCalledWith("test-step", stepFn, options);
   });
 
+  test("should have hasRunningOperations method that returns false initially", () => {
+    const durableContext = createDurableContext(
+      mockExecutionContext,
+      mockParentContext,
+    );
+
+    expect(durableContext.hasRunningOperations()).toBe(false);
+  });
+
   test("should call block handler when runInChildContext method is invoked", () => {
     const durableContext = createDurableContext(
       mockExecutionContext,

packages/lambda-durable-functions-sdk-js/src/context/durable-context/durable-context.ts

@@ -53,13 +53,27 @@ export const createDurableContext = (
   );
 
   let stepCounter = 0;
+  const runningOperations = new Set<string>();
   const checkpoint = createCheckpoint(executionContext, checkpointToken || "");
 
   const createStepId = (): string => {
     stepCounter++;
     return stepPrefix ? `${stepPrefix}-${stepCounter}` : `${stepCounter}`;
   };
 
+  // Internal helpers for managing running operations
+  const addRunningOperation = (stepId: string): void => {
+    runningOperations.add(stepId);
+  };
+
+  const removeRunningOperation = (stepId: string): void => {
+    runningOperations.delete(stepId);
+  };
+
+  const hasRunningOperations = (): boolean => {
+    return runningOperations.size > 0;
+  };
+
   const step: DurableContext["step"] = <T>(
     nameOrFn: string | undefined | StepFunc<T>,
     fnOrOptions?: StepFunc<T> | StepConfig<T>,
@@ -71,6 +85,8 @@ export const createDurableContext = (
       parentContext,
       createStepId,
       createContextLogger,
+      addRunningOperation,
+      removeRunningOperation,
     );
     return stepHandler(nameOrFn, fnOrOptions, maybeOptions);
   };
@@ -177,6 +193,7 @@ export const createDurableContext = (
     ...parentContext,
     _stepPrefix: stepPrefix,
     _stepCounter: stepCounter,
+    hasRunningOperations,
     step,
     runInChildContext,
     wait: createWaitHandler(executionContext, checkpoint, createStepId),
@@ -185,6 +202,8 @@ export const createDurableContext = (
       checkpoint,
       createStepId,
       createContextLogger,
+      addRunningOperation,
+      removeRunningOperation,
     ),
     createCallback,
     waitForCallback,

packages/lambda-durable-functions-sdk-js/src/handlers/step-handler/step-handler.test.ts

@@ -83,6 +83,8 @@ describe("Step Handler", () => {
       mockParentContext,
       createStepId,
       createMockEnrichedLogger,
+      jest.fn(), // addRunningOperation
+      jest.fn(), // removeRunningOperation
     );
 
     // Reset the mock for retryPresets.default

packages/lambda-durable-functions-sdk-js/src/handlers/step-handler/step-handler.ts

@@ -48,6 +48,8 @@ export const createStepHandler = (
   parentContext: Context,
   createStepId: () => string,
   createContextLogger: (stepId: string, attempt?: number) => Logger,
+  addRunningOperation: (stepId: string) => void,
+  removeRunningOperation: (stepId: string) => void,
 ) => {
   return async <T>(
     nameOrFn: string | undefined | StepFunc<T>,
@@ -166,6 +168,8 @@ export const createStepHandler = (
       name,
       fn,
       createContextLogger,
+      addRunningOperation,
+      removeRunningOperation,
       options,
     );
   };
@@ -205,6 +209,8 @@ export const executeStep = async <T>(
   name: string | undefined,
   fn: StepFunc<T>,
   createContextLogger: (stepId: string, attempt?: number) => Logger,
+  addRunningOperation: (stepId: string) => void,
+  removeRunningOperation: (stepId: string) => void,
   options?: StepConfig<T>,
 ): Promise<T> => {
   // Determine step semantics (default to AT_LEAST_ONCE_PER_RETRY if not specified)
@@ -248,9 +254,15 @@ export const executeStep = async <T>(
     };
 
     // Execute the step function with stepContext
-    const result = await OperationInterceptor.forExecution(
-      context.durableExecutionArn,
-    ).execute(name, () => fn(stepContext));
+    addRunningOperation(stepId);
+    let result: T;
+    try {
+      result = await OperationInterceptor.forExecution(
+        context.durableExecutionArn,
+      ).execute(name, () => fn(stepContext));
+    } finally {
+      removeRunningOperation(stepId);
+    }
 
     // Serialize the result for consistency
     const serializedResult = await safeSerialize(

packages/lambda-durable-functions-sdk-js/src/handlers/wait-for-condition-handler/wait-for-condition-handler.test.ts

@@ -81,6 +81,8 @@ describe("WaitForCondition Handler", () => {
       mockCheckpoint,
       createStepId,
       createMockEnrichedLogger,
+      jest.fn(), // addRunningOperation
+      jest.fn(), // removeRunningOperation
     );
   });
 

packages/lambda-durable-functions-sdk-js/src/handlers/wait-for-condition-handler/wait-for-condition-handler.ts

@@ -43,6 +43,8 @@ export const createWaitForConditionHandler = (
   checkpoint: ReturnType<typeof createCheckpoint>,
   createStepId: () => string,
   createContextLogger: (stepId: string, attempt?: number) => Logger,
+  addRunningOperation: (stepId: string) => void,
+  removeRunningOperation: (stepId: string) => void,
 ) => {
   return async <T>(
     nameOrCheck: string | undefined | WaitForConditionCheckFunc<T>,
@@ -111,6 +113,8 @@ export const createWaitForConditionHandler = (
       check,
       config,
       createContextLogger,
+      addRunningOperation,
+      removeRunningOperation,
     );
   };
 };
@@ -150,6 +154,8 @@ export const executeWaitForCondition = async <T>(
   check: WaitForConditionCheckFunc<T>,
   config: WaitForConditionConfig<T>,
   createContextLogger: (stepId: string, attempt?: number) => Logger,
+  addRunningOperation: (stepId: string) => void,
+  removeRunningOperation: (stepId: string) => void,
 ): Promise<T> => {
   const serdes = config.serdes || defaultSerdes;
 
@@ -215,9 +221,15 @@ export const executeWaitForCondition = async <T>(
     };
 
     // Execute the check function
-    const newState = await OperationInterceptor.forExecution(
-      context.durableExecutionArn,
-    ).execute(name, () => check(currentState, waitForConditionContext));
+    addRunningOperation(stepId);
+    let newState: T;
+    try {
+      newState = await OperationInterceptor.forExecution(
+        context.durableExecutionArn,
+      ).execute(name, () => check(currentState, waitForConditionContext));
+    } finally {
+      removeRunningOperation(stepId);
+    }
 
     // Serialize the new state for consistency
     const serializedState = await safeSerialize(

packages/lambda-durable-functions-sdk-js/src/types/index.ts

@@ -87,6 +87,7 @@ export type DurableExecutionInvocationOutput =
 export interface DurableContext extends Context {
   _stepPrefix?: string;
   _stepCounter: number;
+  hasRunningOperations(): boolean;
   step: <T>(
     nameOrFn: string | undefined | StepFunc<T>,
     fnOrOptions?: StepFunc<T> | StepConfig<T>,