Use callback-reported close time in CHASM Nexus completion events

chasm/lib/nexusoperation/operation.go

@@ -71,10 +71,10 @@ const TaskGroupName = "nexus"
 // It's the responsibility of the parent component to apply the appropriate state transitions to the operation.
 type OperationStore interface {
 	OnNexusOperationStarted(ctx chasm.MutableContext, operation *Operation, operationToken string, startTime *time.Time, links []*commonpb.Link) error
-	OnNexusOperationCanceled(ctx chasm.MutableContext, operation *Operation, cause *failurepb.Failure) error
-	OnNexusOperationFailed(ctx chasm.MutableContext, operation *Operation, cause *failurepb.Failure) error
+	OnNexusOperationCanceled(ctx chasm.MutableContext, operation *Operation, cause *failurepb.Failure, closeTime *time.Time) error
+	OnNexusOperationFailed(ctx chasm.MutableContext, operation *Operation, cause *failurepb.Failure, closeTime *time.Time) error
 	OnNexusOperationTimedOut(ctx chasm.MutableContext, operation *Operation, cause *failurepb.Failure, fromAttempt bool) error
-	OnNexusOperationCompleted(ctx chasm.MutableContext, operation *Operation, result *commonpb.Payload, links []*commonpb.Link) error
+	OnNexusOperationCompleted(ctx chasm.MutableContext, operation *Operation, result *commonpb.Payload, closeTime *time.Time, links []*commonpb.Link) error
 	OnNexusOperationCancellationCompleted(ctx chasm.MutableContext, operation *Operation) error
 	OnNexusOperationCancellationFailed(ctx chasm.MutableContext, operation *Operation, cause *failurepb.Failure) error
 	// NexusOperationInvocationData loads invocation data (Input, Header, NexusLinks) from the scheduled history event.
@@ -212,28 +212,31 @@ func (o *Operation) onStarted(ctx chasm.MutableContext, operationToken string, s
 }
 
 // onCompleted applies the succeeded transition or delegates to the store if one is present.
-func (o *Operation) onCompleted(ctx chasm.MutableContext, result *commonpb.Payload, links []*commonpb.Link) error {
+// closeTime, when non-nil, is the callback-reported completion time used instead of the current time.
+func (o *Operation) onCompleted(ctx chasm.MutableContext, result *commonpb.Payload, closeTime *time.Time, links []*commonpb.Link) error {
 	if store, ok := o.Store.TryGet(ctx); ok {
-		return store.OnNexusOperationCompleted(ctx, o, result, links)
+		return store.OnNexusOperationCompleted(ctx, o, result, closeTime, links)
 	}
 	o.Links = append(o.Links, links...)
-	return TransitionSucceeded.Apply(o, ctx, EventSucceeded{Result: result})
+	return TransitionSucceeded.Apply(o, ctx, EventSucceeded{Result: result, CompleteTime: closeTime})
 }
 
 // onFailed applies the failed transition or delegates to the store if one is present.
-func (o *Operation) onFailed(ctx chasm.MutableContext, cause *failurepb.Failure) error {
+// closeTime, when non-nil, is the callback-reported completion time used instead of the current time.
+func (o *Operation) onFailed(ctx chasm.MutableContext, cause *failurepb.Failure, closeTime *time.Time) error {
 	if store, ok := o.Store.TryGet(ctx); ok {
-		return store.OnNexusOperationFailed(ctx, o, cause)
+		return store.OnNexusOperationFailed(ctx, o, cause, closeTime)
 	}
-	return TransitionFailed.Apply(o, ctx, EventFailed{Failure: cause})
+	return TransitionFailed.Apply(o, ctx, EventFailed{Failure: cause, CompleteTime: closeTime})
 }
 
 // onCanceled applies the canceled transition or delegates to the store if one is present.
-func (o *Operation) onCanceled(ctx chasm.MutableContext, cause *failurepb.Failure) error {
+// closeTime, when non-nil, is the callback-reported completion time used instead of the current time.
+func (o *Operation) onCanceled(ctx chasm.MutableContext, cause *failurepb.Failure, closeTime *time.Time) error {
 	if store, ok := o.Store.TryGet(ctx); ok {
-		return store.OnNexusOperationCanceled(ctx, o, cause)
+		return store.OnNexusOperationCanceled(ctx, o, cause, closeTime)
 	}
-	return TransitionCanceled.Apply(o, ctx, EventCanceled{Failure: cause})
+	return TransitionCanceled.Apply(o, ctx, EventCanceled{Failure: cause, CompleteTime: closeTime})
 }
 
 // onTimedOut applies the timed out transition or delegates to the store if one is present.
@@ -269,14 +272,15 @@ func (o *Operation) HandleNexusCompletion(
 		links = nil
 	}
 
+	closeTime := timestamp.TimeValuePtr(completion.GetCloseTime())
 	switch outcome := completion.Outcome.(type) {
 	case *persistencespb.ChasmNexusCompletion_Success:
-		return o.onCompleted(ctx, outcome.Success, nil)
+		return o.onCompleted(ctx, outcome.Success, closeTime, nil)
 	case *persistencespb.ChasmNexusCompletion_Failure:
 		if outcome.Failure.GetCanceledFailureInfo() != nil {
-			return o.onCanceled(ctx, outcome.Failure)
+			return o.onCanceled(ctx, outcome.Failure, closeTime)
 		}
-		return o.onFailed(ctx, outcome.Failure)
+		return o.onFailed(ctx, outcome.Failure, closeTime)
 	default:
 		return serviceerror.NewInvalidArgument("invalid completion outcome")
 	}
@@ -352,11 +356,11 @@ func (o *Operation) saveInvocationResult(
 			// HandleNexusCompletion will apply its outcome from the completion callback.
 			return nil, o.onStarted(ctx, r.response.Pending.Token, nil, links)
 		}
-		return nil, o.onCompleted(ctx, r.response.Successful, links)
+		return nil, o.onCompleted(ctx, r.response.Successful, nil, links)
 	case invocationResultCancel:
-		return nil, o.onCanceled(ctx, r.failure)
+		return nil, o.onCanceled(ctx, r.failure, nil)
 	case invocationResultFail:
-		return nil, o.onFailed(ctx, r.failure)
+		return nil, o.onFailed(ctx, r.failure, nil)
 	case invocationResultTimeout:
 		return nil, o.onTimedOut(ctx, r.failure, true)
 	case invocationResultRetry:

chasm/lib/nexusoperation/operation_test.go

@@ -142,14 +142,33 @@ func TestHandleNexusCompletion(t *testing.T) {
 		t.Run("AfterStarted", func(t *testing.T) {
 			ctx := newCtx()
 			op := newStartedOp(t, ctx)
+			closeTime := defaultTime.Add(time.Minute)
 			err := op.HandleNexusCompletion(ctx, &persistencespb.ChasmNexusCompletion{
 				RequestId: op.GetRequestId(),
+				CloseTime: timestamppb.New(closeTime),
 				Outcome: &persistencespb.ChasmNexusCompletion_Success{
 					Success: mustToPayload(t, "result"),
 				},
 			})
 			require.NoError(t, err)
 			require.Equal(t, nexusoperationpb.OPERATION_STATUS_SUCCEEDED, op.GetStatus())
+			// The callback-reported completion time is stamped onto the operation.
+			require.Equal(t, closeTime, op.GetClosedTime().AsTime())
+		})
+
+		t.Run("AfterStartedWithoutCloseTime", func(t *testing.T) {
+			ctx := newCtx()
+			op := newStartedOp(t, ctx)
+			err := op.HandleNexusCompletion(ctx, &persistencespb.ChasmNexusCompletion{
+				RequestId: op.GetRequestId(),
+				Outcome: &persistencespb.ChasmNexusCompletion_Success{
+					Success: mustToPayload(t, "result"),
+				},
+			})
+			require.NoError(t, err)
+			require.Equal(t, nexusoperationpb.OPERATION_STATUS_SUCCEEDED, op.GetStatus())
+			// Falls back to the current component time when no close time is provided.
+			require.Equal(t, defaultTime, op.GetClosedTime().AsTime())
 		})
 
 		t.Run("CompletionBeforeStart", func(t *testing.T) {
@@ -191,14 +210,17 @@ func TestHandleNexusCompletion(t *testing.T) {
 		t.Run("AfterStarted", func(t *testing.T) {
 			ctx := newCtx()
 			op := newStartedOp(t, ctx)
+			closeTime := defaultTime.Add(time.Minute)
 			err := op.HandleNexusCompletion(ctx, &persistencespb.ChasmNexusCompletion{
 				RequestId: op.GetRequestId(),
+				CloseTime: timestamppb.New(closeTime),
 				Outcome: &persistencespb.ChasmNexusCompletion_Failure{
 					Failure: &failurepb.Failure{Message: "oops"},
 				},
 			})
 			require.NoError(t, err)
 			require.Equal(t, nexusoperationpb.OPERATION_STATUS_FAILED, op.GetStatus())
+			require.Equal(t, closeTime, op.GetClosedTime().AsTime())
 		})
 
 		t.Run("CompletionBeforeStart", func(t *testing.T) {
@@ -224,8 +246,10 @@ func TestHandleNexusCompletion(t *testing.T) {
 		t.Run("AfterStarted", func(t *testing.T) {
 			ctx := newCtx()
 			op := newStartedOp(t, ctx)
+			closeTime := defaultTime.Add(time.Minute)
 			err := op.HandleNexusCompletion(ctx, &persistencespb.ChasmNexusCompletion{
 				RequestId: op.GetRequestId(),
+				CloseTime: timestamppb.New(closeTime),
 				Outcome: &persistencespb.ChasmNexusCompletion_Failure{
 					Failure: &failurepb.Failure{
 						Message: "canceled",
@@ -237,6 +261,7 @@ func TestHandleNexusCompletion(t *testing.T) {
 			})
 			require.NoError(t, err)
 			require.Equal(t, nexusoperationpb.OPERATION_STATUS_CANCELED, op.GetStatus())
+			require.Equal(t, closeTime, op.GetClosedTime().AsTime())
 		})
 
 		t.Run("CompletionBeforeStart", func(t *testing.T) {

chasm/lib/nexusoperation/tasks_test_helpers_test.go

@@ -51,6 +51,7 @@ type mockStoreComponent struct {
 	startLinks      []*commonpb.Link
 	completionLinks []*commonpb.Link
 	startTime       *time.Time
+	closeTime       *time.Time
 }
 
 func (m *mockStoreComponent) LifecycleState(_ chasm.Context) chasm.LifecycleState {
@@ -78,17 +79,20 @@ func (m *mockStoreComponent) OnNexusOperationStarted(ctx chasm.MutableContext, o
 	})
 }
 
-func (m *mockStoreComponent) OnNexusOperationCompleted(ctx chasm.MutableContext, op *Operation, result *commonpb.Payload, links []*commonpb.Link) error {
+func (m *mockStoreComponent) OnNexusOperationCompleted(ctx chasm.MutableContext, op *Operation, result *commonpb.Payload, closeTime *time.Time, links []*commonpb.Link) error {
 	m.completionLinks = links
-	return TransitionSucceeded.Apply(op, ctx, EventSucceeded{Result: result})
+	m.closeTime = closeTime
+	return TransitionSucceeded.Apply(op, ctx, EventSucceeded{Result: result, CompleteTime: closeTime})
 }
 
-func (m *mockStoreComponent) OnNexusOperationFailed(ctx chasm.MutableContext, op *Operation, cause *failurepb.Failure) error {
-	return TransitionFailed.Apply(op, ctx, EventFailed{Failure: cause})
+func (m *mockStoreComponent) OnNexusOperationFailed(ctx chasm.MutableContext, op *Operation, cause *failurepb.Failure, closeTime *time.Time) error {
+	m.closeTime = closeTime
+	return TransitionFailed.Apply(op, ctx, EventFailed{Failure: cause, CompleteTime: closeTime})
 }
 
-func (m *mockStoreComponent) OnNexusOperationCanceled(ctx chasm.MutableContext, op *Operation, cause *failurepb.Failure) error {
-	return TransitionCanceled.Apply(op, ctx, EventCanceled{Failure: cause})
+func (m *mockStoreComponent) OnNexusOperationCanceled(ctx chasm.MutableContext, op *Operation, cause *failurepb.Failure, closeTime *time.Time) error {
+	m.closeTime = closeTime
+	return TransitionCanceled.Apply(op, ctx, EventCanceled{Failure: cause, CompleteTime: closeTime})
 }
 
 func (m *mockStoreComponent) OnNexusOperationTimedOut(ctx chasm.MutableContext, op *Operation, cause *failurepb.Failure, fromAttempt bool) error {

chasm/lib/workflow/nexus_methods.go

@@ -80,6 +80,7 @@ func (w *Workflow) OnNexusOperationCanceled(
 	ctx chasm.MutableContext,
 	op *nexusoperation.Operation,
 	cause *failurepb.Failure,
+	closeTime *time.Time,
 ) error {
 	parentData := &chasmworkflowpb.NexusOperationParentData{}
 	if err := op.GetParentData().UnmarshalTo(parentData); err != nil {
@@ -95,6 +96,10 @@ func (w *Workflow) OnNexusOperationCanceled(
 				Failure:          createNexusOperationFailure(op, scheduledEventID, cause),
 			},
 		}
+		if closeTime != nil {
+			// Use the callback-reported completion time instead of the current time.
+			e.EventTime = timestamppb.New(*closeTime)
+		}
 	})
 	return err
 }
@@ -105,6 +110,7 @@ func (w *Workflow) OnNexusOperationFailed(
 	ctx chasm.MutableContext,
 	op *nexusoperation.Operation,
 	cause *failurepb.Failure,
+	closeTime *time.Time,
 ) error {
 	parentData := &chasmworkflowpb.NexusOperationParentData{}
 	if err := op.GetParentData().UnmarshalTo(parentData); err != nil {
@@ -120,6 +126,10 @@ func (w *Workflow) OnNexusOperationFailed(
 				Failure:          createNexusOperationFailure(op, scheduledEventID, cause),
 			},
 		}
+		if closeTime != nil {
+			// Use the callback-reported completion time instead of the current time.
+			e.EventTime = timestamppb.New(*closeTime)
+		}
 	})
 	return err
 }
@@ -130,6 +140,7 @@ func (w *Workflow) OnNexusOperationCompleted(
 	ctx chasm.MutableContext,
 	op *nexusoperation.Operation,
 	result *commonpb.Payload,
+	closeTime *time.Time,
 	links []*commonpb.Link,
 ) error {
 	parentData := &chasmworkflowpb.NexusOperationParentData{}
@@ -146,6 +157,10 @@ func (w *Workflow) OnNexusOperationCompleted(
 			},
 		}
 		e.Links = links
+		if closeTime != nil {
+			// Use the callback-reported completion time instead of the current time.
+			e.EventTime = timestamppb.New(*closeTime)
+		}
 	})
 	return err
 }

chasm/lib/workflow/nexus_methods_test.go

@@ -0,0 +1,111 @@
+package workflow
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	commandpb "go.temporal.io/api/command/v1"
+	enumspb "go.temporal.io/api/enums/v1"
+	failurepb "go.temporal.io/api/failure/v1"
+	historypb "go.temporal.io/api/history/v1"
+	"go.temporal.io/server/chasm/lib/nexusoperation"
+	"google.golang.org/protobuf/types/known/timestamppb"
+)
+
+// scheduleTestNexusOperation schedules a Nexus operation on the test workflow and returns the
+// scheduled operation together with its scheduled event ID.
+func scheduleTestNexusOperation(t *testing.T, tcx testContext) *nexusoperation.Operation {
+	t.Helper()
+	err := tcx.scheduleHandler(tcx.chasmCtx, tcx.wf, commandValidator{maxPayloadSize: 1}, &commandpb.Command{
+		Attributes: &commandpb.Command_ScheduleNexusOperationCommandAttributes{
+			ScheduleNexusOperationCommandAttributes: &commandpb.ScheduleNexusOperationCommandAttributes{
+				Endpoint:  "endpoint",
+				Service:   "service",
+				Operation: "op",
+			},
+		},
+	}, CommandHandlerOptions{WorkflowTaskCompletedEventID: 1})
+	require.NoError(t, err)
+	require.Len(t, tcx.history.Events, 1)
+
+	event := tcx.history.Events[0]
+	opField, ok := tcx.wf.Operations[event.EventId]
+	require.True(t, ok)
+	return opField.Get(tcx.chasmCtx)
+}
+
+// TestOnNexusOperationCompletion_UsesCloseTime verifies that the completion/failure/cancellation
+// history events stamp the callback-reported close time when provided, and fall back to the
+// current time otherwise.
+func TestOnNexusOperationCompletion_UsesCloseTime(t *testing.T) {
+	closeTime := time.Date(2020, 6, 1, 12, 0, 0, 0, time.UTC)
+	fixedNow := time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
+
+	failure := &failurepb.Failure{Message: "oops"}
+
+	for _, tc := range []struct {
+		name   string
+		invoke func(tcx testContext, op *nexusoperation.Operation, closeTime *time.Time) error
+	}{
+		{
+			name: "Completed",
+			invoke: func(tcx testContext, op *nexusoperation.Operation, closeTime *time.Time) error {
+				return tcx.wf.OnNexusOperationCompleted(tcx.chasmCtx, op, nil, closeTime, nil)
+			},
+		},
+		{
+			name: "Failed",
+			invoke: func(tcx testContext, op *nexusoperation.Operation, closeTime *time.Time) error {
+				return tcx.wf.OnNexusOperationFailed(tcx.chasmCtx, op, failure, closeTime)
+			},
+		},
+		{
+			name: "Canceled",
+			invoke: func(tcx testContext, op *nexusoperation.Operation, closeTime *time.Time) error {
+				return tcx.wf.OnNexusOperationCanceled(tcx.chasmCtx, op, failure, closeTime)
+			},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Run("uses close time", func(t *testing.T) {
+				tcx := newTestContext(t, defaultConfig)
+				op := scheduleTestNexusOperation(t, tcx)
+				overrideDefaultEventTime(tcx, fixedNow)
+
+				require.NoError(t, tc.invoke(tcx, op, &closeTime))
+
+				completionEvent := tcx.history.Events[len(tcx.history.Events)-1]
+				require.Equal(t, closeTime, completionEvent.GetEventTime().AsTime())
+			})
+
+			t.Run("falls back to current time", func(t *testing.T) {
+				tcx := newTestContext(t, defaultConfig)
+				op := scheduleTestNexusOperation(t, tcx)
+				overrideDefaultEventTime(tcx, fixedNow)
+
+				require.NoError(t, tc.invoke(tcx, op, nil))
+
+				completionEvent := tcx.history.Events[len(tcx.history.Events)-1]
+				require.Equal(t, fixedNow, completionEvent.GetEventTime().AsTime())
+			})
+		})
+	}
+}
+
+// overrideDefaultEventTime pins the default EventTime assigned by the backend to newly added
+// history events, so fallback-to-current-time behavior can be asserted deterministically.
+func overrideDefaultEventTime(tcx testContext, now time.Time) {
+	nextEventID := int64(len(tcx.history.Events) + 5)
+	tcx.backend.HandleAddHistoryEvent = func(eventType enumspb.EventType, setAttrs func(*historypb.HistoryEvent)) *historypb.HistoryEvent {
+		e := &historypb.HistoryEvent{
+			Version:   1,
+			EventId:   nextEventID,
+			EventTime: timestamppb.New(now),
+		}
+		nextEventID++
+		setAttrs(e)
+		tcx.history.Events = append(tcx.history.Events, e)
+		return e
+	}
+}

tests/nexus_workflow_test.go

@@ -1073,11 +1073,13 @@ func (s *NexusWorkflowTestSuite) TestNexusOperationAsyncCompletion(chasmEnabled
 	// Send an invalid completion request and verify that we get an error that the namespace in the URL doesn't match the namespace in the token.
 	invalidCallbackURL := "http://" + env.HttpAPIAddress() + "/" + commonnexus.RouteCompletionCallback.Path(invalidNamespace)
 
+	closeTime := time.Now().Truncate(time.Millisecond).UTC()
 	completion := nexusrpc.CompleteOperationOptions{
 		Result: testcore.MustToPayload(s.T(), "result"),
 		Header: nexus.Header{commonnexus.CallbackTokenHeader: callbackToken},
 		// Repeat the handler link on completion to verify callback links do not leak onto the completed event.
-		Links: []nexus.Link{handlerNexusLink},
+		Links:     []nexus.Link{handlerNexusLink},
+		CloseTime: closeTime,
 	}
 	err = s.sendNexusCompletionRequest(ctx, invalidCallbackURL, completion)
 	// Verify we get the correct error response
@@ -1228,6 +1230,16 @@ func (s *NexusWorkflowTestSuite) TestNexusOperationAsyncCompletion(chasmEnabled
 	s.Len(opStartedEvent.Links, 1)
 	protorequire.ProtoEqual(s.T(), handlerLink, opStartedEvent.Links[0].GetWorkflowEvent())
 
+	if chasmEnabled {
+		// The completed event must be stamped with the callback-reported close time, not the server
+		// time at which the completion was processed. The HSM path does not carry CloseTime, so this
+		// assertion is CHASM-only.
+		completedEvent := s.RequireHistoryEvent(hist, enumspb.EVENT_TYPE_NEXUS_OPERATION_COMPLETED)
+		s.True(closeTime.Equal(completedEvent.GetEventTime().AsTime()),
+			"completed event must carry the callback close time; got %s want %s",
+			completedEvent.GetEventTime().AsTime(), closeTime)
+	}
+
 	// Find the first WFT completed after the started event for reset tests.
 	wftCompletedIdx := slices.IndexFunc(hist, func(e *historypb.HistoryEvent) bool {
 		return e.EventType == enumspb.EVENT_TYPE_WORKFLOW_TASK_COMPLETED && e.EventId > opStartedEvent.EventId