add more comments and remove registration from context

Author: zhongkechen

sdk/src/main/java/software/amazon/lambda/durable/BaseContext.java

@@ -14,7 +14,7 @@ public abstract class BaseContext {
     private final String contextId;
     private boolean isReplaying;
 
-    /** Shared initialization — sets all fields but performs no thread registration. */
+    /** Creates a new BaseContext instance. */
     protected BaseContext(
             ExecutionManager executionManager, DurableConfig durableConfig, Context lambdaContext, String contextId) {
         this.executionManager = executionManager;

sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java

@@ -10,8 +10,6 @@
 import java.util.function.Supplier;
 import org.slf4j.LoggerFactory;
 import software.amazon.lambda.durable.execution.ExecutionManager;
-import software.amazon.lambda.durable.execution.ThreadContext;
-import software.amazon.lambda.durable.execution.ThreadType;
 import software.amazon.lambda.durable.logging.DurableLogger;
 import software.amazon.lambda.durable.operation.CallbackOperation;
 import software.amazon.lambda.durable.operation.ChildContextOperation;
@@ -24,7 +22,7 @@ public class DurableContext extends BaseContext {
     private final AtomicInteger operationCounter;
     private final DurableLogger logger;
 
-    /** Shared initialization — sets all fields but performs no thread registration. */
+    /** Shared initialization — sets all fields. */
     private DurableContext(
             ExecutionManager executionManager, DurableConfig durableConfig, Context lambdaContext, String contextId) {
         super(executionManager, durableConfig, lambdaContext, contextId);
@@ -39,7 +37,7 @@ private DurableContext(
     }
 
     /**
-     * Creates a root context and registers the current thread for execution coordination.
+     * Creates a root context (contextId = null)
      *
      * <p>The context itself always has a null contextId (making it a root context).
      *
@@ -50,16 +48,11 @@ private DurableContext(
      */
     public static DurableContext createRootContext(
             ExecutionManager executionManager, DurableConfig durableConfig, Context lambdaContext) {
-        var ctx = new DurableContext(executionManager, durableConfig, lambdaContext, null);
-        executionManager.registerActiveThread(null);
-        executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT));
-        return ctx;
+        return new DurableContext(executionManager, durableConfig, lambdaContext, null);
     }
 
     /**
-     * Creates a child context without registering the current thread. Thread registration is handled by
-     * ChildContextOperation, which registers on the parent thread before the executor runs and sets the context on the
-     * child thread inside the executor.
+     * Creates a child context.
      *
      * @param childContextId the child context's ID (the CONTEXT operation's operation ID)
      * @return a new DurableContext for the child context

sdk/src/main/java/software/amazon/lambda/durable/DurableExecutor.java

@@ -19,6 +19,8 @@
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.execution.ExecutionManager;
 import software.amazon.lambda.durable.execution.SuspendExecutionException;
+import software.amazon.lambda.durable.execution.ThreadContext;
+import software.amazon.lambda.durable.execution.ThreadType;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
 import software.amazon.lambda.durable.model.DurableExecutionOutput;
 import software.amazon.lambda.durable.serde.SerDes;
@@ -45,6 +47,8 @@ public static <I, O> DurableExecutionOutput execute(
                             extractUserInput(executionManager.getExecutionOperation(), config.getSerDes(), inputType);
                     // Create context in the executor thread so it detects the correct thread name
                     var context = DurableContext.createRootContext(executionManager, config, lambdaContext);
+                    executionManager.registerActiveThread(null);
+                    executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT));
                     return handler.apply(userInput, context);
                 },
                 config.getExecutorService()); // Get executor from config for running user code

sdk/src/main/java/software/amazon/lambda/durable/StepContext.java

@@ -11,12 +11,12 @@ public class StepContext extends BaseContext {
     private final DurableLogger logger;
 
     /**
-     * Shared initialization — sets all fields but performs no thread registration.
+     * Creates a new StepContext instance for use in step operations.
      *
      * @param executionManager Manages durable execution state and operations
      * @param durableConfig Configuration for durable execution behavior
      * @param lambdaContext AWS Lambda runtime context
-     * @param stepOperationId Unique identifier for this context instance
+     * @param stepOperationId Unique identifier for this context instance that equals to step operation id
      */
     protected StepContext(
             ExecutionManager executionManager,