CHASM return same timestamp from context.Now() within same MutableCtx

chasm/context.go

@@ -98,6 +98,8 @@ type immutableCtx struct {
 
 type mutableCtx struct {
 	*immutableCtx
+
+	nowByComponent map[*Node]time.Time
 }
 
 // NewContext creates a new Context from an existing Context and root Node.
@@ -209,8 +211,27 @@ func NewMutableContext(
 	node *Node,
 ) MutableContext {
 	return &mutableCtx{
-		immutableCtx: newContext(ctx, node),
+		immutableCtx:   newContext(ctx, node),
+		nowByComponent: make(map[*Node]time.Time),
+	}
+}
+
+func (c *mutableCtx) Now(component Component) time.Time {
+	node, ok := c.root.valueToNode[component]
+	if !ok || !node.isComponent() {
+		return c.root.Now(component)
+	}
+
+	if now, ok := c.nowByComponent[node]; ok {
+		return now
+	}
+
+	now := c.root.Now(component)
+	if c.nowByComponent == nil {
+		c.nowByComponent = make(map[*Node]time.Time)
 	}
+	c.nowByComponent[node] = now
+	return now
 }
 
 func (c *mutableCtx) AddTask(
@@ -223,7 +244,8 @@ func (c *mutableCtx) AddTask(
 
 func (c *mutableCtx) withValue(key any, value any) Context {
 	return &mutableCtx{
-		immutableCtx: ContextWithValue(c.immutableCtx, key, value),
+		immutableCtx:   ContextWithValue(c.immutableCtx, key, value),
+		nowByComponent: c.nowByComponent,
 	}
 }
 

chasm/tree_test.go

@@ -3127,6 +3127,37 @@
 	return node
 }
 
+func (s *nodeSuite) TestMutableContextNowStableWithinContext() {
+	root := s.testComponentTree()
+
+	startTime := time.Date(2026, 1, 1, 1, 0, 0, 0, time.UTC)
+	updatedTime := startTime.Add(time.Minute)
+	laterTime := updatedTime.Add(time.Minute)
+
+	s.timeSource.Update(startTime)
+
+	mutableContext := NewMutableContext(context.Background(), root)
+	component, err := root.Component(mutableContext, ComponentRef{})
+	s.NoError(err)
+
+	s.Equal(startTime, mutableContext.Now(component))
+
+	s.timeSource.Update(updatedTime)
+	s.Equal(startTime, mutableContext.Now(component))
+
+	contextWithValue := ContextWithValue(mutableContext, "test-key", "test-value")
+	s.Equal(startTime, contextWithValue.Now(component))
+
+	newMutableContext := NewMutableContext(context.Background(), root)
+	s.Equal(updatedTime, newMutableContext.Now(component))
+
+	immutableContext := NewContext(context.Background(), root)
+	s.Equal(updatedTime, immutableContext.Now(component))
+
+	s.timeSource.Update(laterTime)
+	s.Equal(laterTime, immutableContext.Now(component))
+}
+
 func (s *nodeSuite) TestExecuteImmediatePureTask() {
 	root := s.testComponentTree()
 
@@ -3182,6 +3213,68 @@
 	s.Equal(tasks.MaximumKey.FireTime, s.nodeBackend.LastDeletePureTaskCall())
 }
 
+func (s *nodeSuite) TestImmediatePureTaskNowStableWithinTaskOnly() {
+	root := s.testComponentTree()
+
+	_, err := root.CloseTransaction()
+	s.NoError(err)
+
+	taskStartTime := time.Date(2026, 1, 1, 2, 0, 0, 0, time.UTC)
+	nextTaskTime := taskStartTime.Add(time.Minute)
+	s.timeSource.Update(taskStartTime)
+
+	mutableContext := NewMutableContext(context.Background(), root)
+	component, err := root.Component(mutableContext, ComponentRef{})
+	s.NoError(err)
+
+	taskAttributes := TaskAttributes{ScheduledTime: TaskScheduledTimeImmediate}
+	mutableContext.AddTask(
+		component,
+		taskAttributes,
+		&TestPureTask{
+			Payload: &commonpb.Payload{Data: []byte("task-1")},
+		},
+	)
+	mutableContext.AddTask(
+		component,
+		taskAttributes,
+		&TestPureTask{
+			Payload: &commonpb.Payload{Data: []byte("task-2")},
+		},
+	)
+
+	s.testLibrary.mockPureTaskHandler.EXPECT().
+		Validate(gomock.Any(), gomock.Any(), gomock.Eq(taskAttributes), gomock.Any()).Return(true, nil).Times(2)
+
+	var observedTimes []time.Time
+	s.testLibrary.mockPureTaskHandler.EXPECT().
+		Execute(
+			gomock.AssignableToTypeOf(&mutableCtx{}),
+			gomock.AssignableToTypeOf(&TestComponent{}),
+			gomock.Eq(taskAttributes),
+			gomock.Any(),
+		).
+		DoAndReturn(func(ctx MutableContext, component any, _ TaskAttributes, _ *TestPureTask) error {
+			chasmComponent := component.(Component)
+			firstNow := ctx.Now(chasmComponent)
+			secondNow := ctx.Now(chasmComponent)
+			s.Equal(firstNow, secondNow)
+
+			observedTimes = append(observedTimes, firstNow)
+			if len(observedTimes) == 1 {
+				s.timeSource.Update(nextTaskTime)
+			}
+			return nil
+		}).
+		Times(2)
+
+	mutations, err := root.CloseTransaction()
+	s.NoError(err)
+	s.Empty(mutations.UpdatedNodes)
+	s.Empty(mutations.DeletedNodes)
+	s.Equal([]time.Time{taskStartTime, nextTaskTime}, observedTimes)
+}
+
 func (s *nodeSuite) TestEachPureTask() {
 	now := s.timeSource.Now()
 
@@ -3221,14 +3314,14 @@
 				InitialVersionedTransition: &persistencespb.VersionedTransition{TransitionCount: 1},
 				Attributes: &persistencespb.ChasmNodeMetadata_ComponentAttributes{
 					ComponentAttributes: &persistencespb.ChasmComponentAttributes{
 						TypeId: testSubComponent1TypeID,
 						PureTasks: []*persistencespb.ChasmComponentAttributes_Task{
 							{
 								TypeId: testPureTaskTypeID,
 								// Not expired yet.
 								ScheduledTime:             timestamppb.New(now.Add(time.Hour)),
 								VersionedTransition:       &persistencespb.VersionedTransition{TransitionCount: 1},
 								VersionedTransitionOffset: 2,
 								PhysicalTaskStatus:        physicalTaskStatusCreated,
 								Data: mustEncode(&commonpb.Payload{
 									Data: []byte("some-random-data-sc1"),