Dispatch activity cancellation to worker using Nexus

common/metrics/metric_defs.go

@@ -748,6 +748,11 @@ var (
 		"nexus_completion_request_preprocess_errors",
 		WithDescription("The number of Nexus completion requests for which pre-processing failed."),
 	)
+	WorkerCommandsSent = NewCounterDef(
+		"worker_commands_sent",
+		WithDescription("The number of worker command dispatches, tagged by outcome (e.g. success, no_poller, rpc_error)."),
+	)
+
 	HostRPSLimit          = NewGaugeDef("host_rps_limit")
 	NamespaceHostRPSLimit = NewGaugeDef("namespace_host_rps_limit")
 	HandoverWaitLatency   = NewTimerDef("handover_wait_latency")

proto/internal/temporal/server/api/persistence/v1/executions.proto

@@ -678,6 +678,21 @@ message ActivityInfo {
   // A dedicated per-worker Nexus task queue on which the server sends control
   // tasks (e.g. activity cancellation) to this specific worker instance.
   string worker_control_task_queue = 51;
+
+  // The shard clock at the time this activity was started (RecordActivityTaskStarted).
+  // Matching uses this clock to build the task token sent to the worker. Stored here so
+  // that history can later reconstruct the same task token (e.g. for cancel worker commands).
+  //
+  // IMPORTANT: The clock approach requires history to reconstruct the token using
+  // the same fields and logic as matching — if NewActivityTaskToken changes, both
+  // call sites must stay in sync or the tokens will silently diverge. An alternative
+  // is to store the full serialized task token (~150-300 bytes), which avoids
+  // reconstruction entirely and is immune to token format changes. We chose the
+  // clock approach to keep the per-activity memory footprint minimal (~24 bytes).
+  //
+  // Replication: This field is part of ActivityInfo and is automatically replicated
+  // via state-based replication. No special handling is needed.
+  temporal.server.api.clock.v1.VectorClock started_clock = 52;
 }
 
 // timer_map column

service/history/api/recordactivitytaskstarted/api.go

@@ -152,6 +152,16 @@ func recordActivityTaskStarted(
 		if ai.RequestId == requestID {
 			response.StartedTime = ai.StartedTime
 			response.Attempt = ai.Attempt
+			if ai.StartedClock != nil {
+				response.Clock = ai.StartedClock
+			} else {
+				// Activity started before the StartedClock field was added.
+				// Create a fresh clock for the shard staleness check.
+				response.Clock, err = shardContext.NewVectorClock()
+				if err != nil {
+					return nil, rejectCodeUndefined, err
+				}
+			}
 			return response, rejectCodeAccepted, nil
 		}
 
@@ -238,6 +248,17 @@ func recordActivityTaskStarted(
 		}
 	}
 
+	// Create the shard clock before recording the start event. Matching uses the returned clock
+	// to build the task token sent to the worker. We store it in ActivityInfo so that history
+	// can later reconstruct the same task token that matching created (e.g. for cancel worker
+	// commands). On retries of this RPC, we return the stored clock from the early return
+	// path above.
+	clock, err := shardContext.NewVectorClock()
+	if err != nil {
+		return nil, rejectCodeUndefined, err
+	}
+	ai.StartedClock = clock
+
 	versioningStamp := worker_versioning.StampFromCapabilities(request.PollRequest.WorkerVersionCapabilities, request.PollRequest.DeploymentOptions) //nolint:staticcheck // SA1019: WorkerVersionCapabilities is deprecated but still used for old versioning [cleanup-old-wv]
 	if _, err := mutableState.AddActivityTaskStartedEvent(
 		ai, scheduledEventID, requestID, request.PollRequest.GetIdentity(),
@@ -261,6 +282,8 @@ func recordActivityTaskStarted(
 		),
 	).Record(scheduleToStartLatency)
 
+	response.Clock = clock
+
 	response.StartedTime = ai.StartedTime
 	response.Attempt = ai.Attempt
 	response.HeartbeatDetails = ai.LastHeartbeatDetails

service/history/api/recordactivitytaskstarted/api_test.go

@@ -4,14 +4,23 @@ import (
 	"context"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	commonpb "go.temporal.io/api/common/v1"
 	enumspb "go.temporal.io/api/enums/v1"
+	historypb "go.temporal.io/api/history/v1"
+	clockspb "go.temporal.io/server/api/clock/v1"
 	deploymentspb "go.temporal.io/server/api/deployment/v1"
+	"go.temporal.io/server/api/historyservice/v1"
 	"go.temporal.io/server/api/matchingservice/v1"
 	"go.temporal.io/server/api/matchingservicemock/v1"
 	persistencespb "go.temporal.io/server/api/persistence/v1"
+	"go.temporal.io/server/common/cluster"
+	"go.temporal.io/server/common/metrics"
+	"go.temporal.io/server/common/namespace"
 	"go.temporal.io/server/common/worker_versioning"
+	historyi "go.temporal.io/server/service/history/interfaces"
 	"go.uber.org/mock/gomock"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -265,3 +274,106 @@ func TestGetDeploymentVersionForWorkflowID_UnversionedTaskQueue(t *testing.T) {
 	assert.Nil(t, targetVersion, "Unversioned task queue should return nil version")
 	assert.Equal(t, int64(0), targetRevNum, "Unversioned task queue should return 0 revision number")
 }
+
+const testClusterName = "active"
+
+// setupMutableStateWithStartedActivity creates mock shard and mutable state with an activity
+// that is already started (StartedEventId != EmptyEventID). The activity's StartedClock is set
+// to the provided value. Returns the mocks and a request whose RequestId matches the activity's,
+// so the call will hit the duplicate-request early-return path in recordActivityTaskStarted.
+func setupMutableStateWithStartedActivity(t *testing.T, startedClock *clockspb.VectorClock) (
+	*historyi.MockShardContext,
+	*historyi.MockMutableState,
+	*historyservice.RecordActivityTaskStartedRequest,
+) {
+	t.Helper()
+	ctrl := gomock.NewController(t)
+
+	mockShard := historyi.NewMockShardContext(ctrl)
+	mockMS := historyi.NewMockMutableState(ctrl)
+
+	nsID := uuid.New().String()
+	scheduledEventID := int64(5)
+	requestID := "test-request-id"
+
+	// Namespace registry
+	nsEntry := namespace.NewLocalNamespaceForTest(
+		&persistencespb.NamespaceInfo{Id: nsID, Name: "test-namespace"},
+		&persistencespb.NamespaceConfig{},
+		testClusterName,
+	)
+	mockNSReg := namespace.NewMockRegistry(ctrl)
+	mockNSReg.EXPECT().GetNamespaceByID(namespace.ID(nsID)).Return(nsEntry, nil)
+
+	// Cluster metadata
+	mockClusterMeta := cluster.NewMockMetadata(ctrl)
+	mockClusterMeta.EXPECT().GetCurrentClusterName().Return(testClusterName)
+
+	mockShard.EXPECT().GetNamespaceRegistry().Return(mockNSReg)
+	mockShard.EXPECT().GetClusterMetadata().Return(mockClusterMeta)
+	mockShard.EXPECT().GetMetricsHandler().Return(metrics.NoopMetricsHandler).AnyTimes()
+
+	ai := &persistencespb.ActivityInfo{
+		ScheduledEventId: scheduledEventID,
+		StartedEventId:   7, // already started
+		RequestId:        requestID,
+		StartedTime:      timestamppb.Now(),
+		Attempt:          1,
+		StartedClock:     startedClock,
+	}
+
+	mockMS.EXPECT().GetActivityInfo(scheduledEventID).Return(ai, true)
+	mockMS.EXPECT().GetActivityScheduledEvent(gomock.Any(), scheduledEventID).Return(
+		&historypb.HistoryEvent{EventId: scheduledEventID}, nil,
+	)
+	mockMS.EXPECT().GetExecutionInfo().Return(&persistencespb.WorkflowExecutionInfo{})
+
+	request := &historyservice.RecordActivityTaskStartedRequest{
+		NamespaceId: nsID,
+		WorkflowExecution: &commonpb.WorkflowExecution{
+			WorkflowId: "test-wf-id",
+			RunId:      "test-run-id",
+		},
+		ScheduledEventId: scheduledEventID,
+		RequestId:        requestID,
+	}
+
+	return mockShard, mockMS, request
+}
+
+// TestRecordActivityTaskStarted_DuplicateRequest_NilStartedClock verifies that
+// when a duplicate RecordActivityTaskStarted request arrives for an activity
+// started before the StartedClock field was added (StartedClock is nil), the response
+// still contains a non-nil Clock for the shard staleness check.
+func TestRecordActivityTaskStarted_DuplicateRequest_NilStartedClock(t *testing.T) {
+	mockShard, mockMS, request := setupMutableStateWithStartedActivity(t, nil /* no StartedClock */)
+
+	// Should call NewVectorClock as fallback for nil StartedClock
+	fallbackClock := &clockspb.VectorClock{ClusterId: 1, ShardId: 1, Clock: 42}
+	mockShard.EXPECT().NewVectorClock().Return(fallbackClock, nil)
+
+	resp, code, err := recordActivityTaskStarted(
+		context.Background(), mockShard, mockMS, request, nil, nil,
+	)
+	require.NoError(t, err)
+	require.Equal(t, rejectCodeAccepted, code)
+	require.NotNil(t, resp.Clock, "Clock must be non-nil even for pre-deploy activities")
+	require.Equal(t, fallbackClock, resp.Clock)
+}
+
+// TestRecordActivityTaskStarted_DuplicateRequest_WithStartedClock verifies that
+// when StartedClock is stored, the stored clock is returned without creating a new one.
+func TestRecordActivityTaskStarted_DuplicateRequest_WithStartedClock(t *testing.T) {
+	storedClock := &clockspb.VectorClock{ClusterId: 1, ShardId: 1, Clock: 100}
+	mockShard, mockMS, request := setupMutableStateWithStartedActivity(t, storedClock)
+
+	// Should NOT call NewVectorClock since StartedClock is available
+	mockShard.EXPECT().NewVectorClock().Times(0)
+
+	resp, code, err := recordActivityTaskStarted(
+		context.Background(), mockShard, mockMS, request, nil, nil,
+	)
+	require.NoError(t, err)
+	require.Equal(t, rejectCodeAccepted, code)
+	require.Equal(t, storedClock, resp.Clock, "Should return the stored StartedClock")
+}

service/history/api/respondworkflowtaskcompleted/workflow_task_completed_handler.go

@@ -682,7 +682,6 @@ func (handler *workflowTaskCompletedHandler) handleCommandRequestCancelActivity(
 	if ai != nil {
 		// If ai is nil, the activity has already been canceled/completed/timedout. The cancel request
 		// will be recorded in the history, but no further action will be taken.
-
 		if ai.StartedEventId == common.EmptyEventID {
 			// We haven't started the activity yet, we can cancel the activity right away and
 			// schedule a workflow task to ensure the workflow makes progress.
@@ -697,37 +696,48 @@ func (handler *workflowTaskCompletedHandler) handleCommandRequestCancelActivity(
 				return nil, err
 			}
 			handler.activityNotStartedCancelled = true
-		} else if ai.StartedEventId != common.EmptyEventID && ai.WorkerControlTaskQueue != "" {
-			// Activity has started and worker supports Nexus control tasks - collect for batched dispatch.
-			taskToken, err := handler.tokenSerializer.Serialize(tasktoken.NewActivityTaskToken(
-				handler.mutableState.GetNamespaceEntry().ID().String(),
-				handler.mutableState.GetWorkflowKey().WorkflowID,
-				handler.mutableState.GetWorkflowKey().RunID,
-				ai.ScheduledEventId,
-				ai.ActivityId,
-				ai.ActivityType.GetName(),
-				ai.Attempt,
-				nil, // Clock not needed for cancel
-				ai.Version,
-				ai.StartVersion,
-				nil,
-			))
-			if err != nil {
-				return nil, err
-			}
-			if handler.pendingWorkerCommandsByControlQueue == nil {
-				handler.pendingWorkerCommandsByControlQueue = make(map[string][]*workerpb.WorkerCommand)
-			}
-			handler.pendingWorkerCommandsByControlQueue[ai.WorkerControlTaskQueue] = append(
-				handler.pendingWorkerCommandsByControlQueue[ai.WorkerControlTaskQueue],
-				&workerpb.WorkerCommand{
-					Type: &workerpb.WorkerCommand_CancelActivity{
-						CancelActivity: &workerpb.CancelActivityCommand{
-							TaskToken: taskToken,
+		} else if ai.WorkerControlTaskQueue != "" {
+			if ai.StartedClock == nil {
+				// StartedClock may be nil for activities started before this feature was deployed.
+				// Skip cancel command; the activity will time out normally.
+				handler.logger.Info("Skipping worker cancel command: activity missing StartedClock (pre-deploy)",
+					tag.WorkflowNamespaceID(handler.mutableState.GetWorkflowKey().NamespaceID),
+					tag.WorkflowID(handler.mutableState.GetWorkflowKey().WorkflowID),
+					tag.WorkflowRunID(handler.mutableState.GetWorkflowKey().RunID),
+					tag.WorkflowScheduledEventID(ai.ScheduledEventId),
+				)
+			} else {
+				// Activity has started and worker supports Nexus control tasks - collect for batched dispatch.
+				taskToken, err := handler.tokenSerializer.Serialize(tasktoken.NewActivityTaskToken(
+					handler.mutableState.GetNamespaceEntry().ID().String(),
+					handler.mutableState.GetWorkflowKey().WorkflowID,
+					handler.mutableState.GetWorkflowKey().RunID,
+					ai.ScheduledEventId,
+					ai.ActivityId,
+					ai.ActivityType.GetName(),
+					ai.Attempt,
+					ai.StartedClock,
+					ai.Version,
+					ai.StartVersion,
+					nil,
+				))
+				if err != nil {
+					return nil, err
+				}
+				if handler.pendingWorkerCommandsByControlQueue == nil {
+					handler.pendingWorkerCommandsByControlQueue = make(map[string][]*workerpb.WorkerCommand)
+				}
+				handler.pendingWorkerCommandsByControlQueue[ai.WorkerControlTaskQueue] = append(
+					handler.pendingWorkerCommandsByControlQueue[ai.WorkerControlTaskQueue],
+					&workerpb.WorkerCommand{
+						Type: &workerpb.WorkerCommand_CancelActivity{
+							CancelActivity: &workerpb.CancelActivityCommand{
+								TaskToken: taskToken,
+							},
 						},
 					},
-				},
-			)
+				)
+			}
 		}
 	}
 	return actCancelReqEvent, nil