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

chasm/context.go

@@ -98,6 +98,8 @@ type immutableCtx struct {
 
 type mutableCtx struct {
 	*immutableCtx
+
+	now **time.Time
 }
 
 // NewContext creates a new Context from an existing Context and root Node.
@@ -208,11 +210,28 @@ func NewMutableContext(
 	ctx context.Context,
 	node *Node,
 ) MutableContext {
+	var now *time.Time
 	return &mutableCtx{
 		immutableCtx: newContext(ctx, node),
+		now:          &now,
 	}
 }
 
+func (c *mutableCtx) Now(component Component) time.Time {
+	if c.now == nil {
+		var now *time.Time
+		c.now = &now
+	}
+
+	if *c.now != nil {
+		return **c.now
+	}
+
+	now := c.root.Now(component)
+	*c.now = &now
+	return now
+}
+
 func (c *mutableCtx) AddTask(
 	component Component,
 	attributes TaskAttributes,
@@ -222,8 +241,14 @@ func (c *mutableCtx) AddTask(
 }
 
 func (c *mutableCtx) withValue(key any, value any) Context {
+	if c.now == nil {
+		var now *time.Time
+		c.now = &now
+	}
+
 	return &mutableCtx{
 		immutableCtx: ContextWithValue(c.immutableCtx, key, value),
+		now:          c.now,
 	}
 }
 

chasm/tree_test.go

@@ -3127,6 +3127,39 @@ func (s *nodeSuite) testComponentTree() *Node {
 	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)
+	testComponent := component.(*TestComponent)
+
+	contextWithValue := ContextWithValue(mutableContext, "test-key", "test-value")
+	s.Equal(startTime, contextWithValue.Now(component))
+
+	s.timeSource.Update(updatedTime)
+	s.Equal(startTime, mutableContext.Now(component))
+
+	childComponent := testComponent.SubComponent1.Get(mutableContext)
+	s.Equal(startTime, mutableContext.Now(childComponent))
+
+	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 +3215,64 @@ func (s *nodeSuite) TestExecuteImmediatePureTask() {
 	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{},
+	)
+	mutableContext.AddTask(
+		component,
+		taskAttributes,
+		&TestPureTask{},
+	)
+
+	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()