Include transient and speculative WFT events in GetWorkflowExecutionHistoryResponse (#9325)

## What changed?

Re-does #9138 which was incidentally merged.

Include transient and speculative WFT events in
`GetWorkflowExecutionHistoryReponse` response, unless UI or CLI made
request.

* Adds `transient_or_speculative_events` back to
`GetMutableStateResponse`
* Reserve `transient_workflow_task` in `HisotryCOntinuation` token
* Add validation helpers
* Add query-compare-query for transient events at request start and end

Re-implements #7732

## Why?
Fix "premature end of stream" errors when workers request history after
cache eviction w/ transient/speculative workflow tasks present. This
adds transient & speculative WFT events in `GetWorkflowExecutionHistory`
(already in `PollWorkflowTask`). Worker cache eviction w/ speculative
workflow tasks causes the expected and actual event counts to be
different. #7732 passed transient events through continuation tokens,
which could become stale during pagination. This PR implements mutable
state querying at both start and end of pagination and compares
transient event IDs to detect if WFT state changed during pagination and
return a retryable error.

## How did you test it?
- [X] built
- [X] run locally and tested manually
- [X] covered by existing tests
- [ ] added new unit test(s)
- [ ] added new functional test(s)

## Potential risks
Same risks from #7732

Author: spkane31

api/historyservice/v1/request_response.pb.go

@@ -49,12 +49,14 @@ a special [in-memory-queue](./in-memory-queue.md) is used for speculative Workfl
 > #### TODO
 > It is important to point out that the `WorkflowTaskScheduled` and `WorkflowTaskStarted` events
 > for transient and speculative Workflow Task are only added to the `PollWorkflowTask` response - and
-> not to the `GetWorkflowExecutionHistory` response. This has an unfortunate consequence: when the 
+> not to the `GetWorkflowExecutionHistory` response. This has an unfortunate consequence: when the
 > worker receives a speculative Workflow Task on a sticky task queue, but the Workflow is already
-> evicted from its cache, it issues a `GetWorkflowExecutionHistory` request, which returns the 
-> history *without* speculative events. This leads to a `premature end of stream` error on the 
+> evicted from its cache, it issues a `GetWorkflowExecutionHistory` request, which returns the
+> history *without* speculative events. This leads to a `premature end of stream` error on the
 > worker side. The worker fails the Workflow Task, clears stickiness, and everything works fine
 > after that - but a failed Workflow Task appears in the history. Fortunately, it doesn't happen often.
+>
+> See PR #9325 for related work on ensuring transient events are not incorrectly returned to CLI/UI clients.
 
 ## Speculative Workflow Task & Workflow Update
 Speculative Workflow Task was introduced to make it possible for Workflow Update to have zero writes

api/token/v1/message.pb.go

@@ -154,6 +154,9 @@ message GetMutableStateResponse {
     string inherited_build_id = 23;
     repeated temporal.server.api.persistence.v1.VersionedTransition transition_history = 24;
     temporal.api.workflow.v1.WorkflowExecutionVersioningInfo versioning_info = 25;
+    // Transient or speculative workflow task events which are not yet persisted in the history.
+    // These events should be appended to the history when it is returned to the worker.
+    temporal.server.api.history.v1.TransientWorkflowTaskInfo transient_or_speculative_tasks = 26;
 }
 
 message PollMutableStateRequest {

docs/architecture/speculative-workflow-task.md

@@ -15,7 +15,7 @@ message HistoryContinuation {
     int64 next_event_id = 3;
     bool is_workflow_running = 5;
     bytes persistence_token = 6;
-    temporal.server.api.history.v1.TransientWorkflowTaskInfo transient_workflow_task = 7;
+    reserved 7; // Was: transient_workflow_task - no longer passed through continuation token
     bytes branch_token = 8;
     reserved 9;
     temporal.server.api.history.v1.VersionHistoryItem version_history_item = 10;

proto/internal/temporal/server/api/historyservice/v1/request_response.proto

@@ -115,24 +115,21 @@ func GetRawHistory(
 	metrics.HistorySize.With(metricsHandler).Record(int64(size))
 
 	if len(nextToken) == 0 && transientWorkflowTaskInfo != nil {
-		if err := validateTransientWorkflowTaskEvents(nextEventID, transientWorkflowTaskInfo); err != nil {
-			logger := shardContext.GetLogger()
-			metricsHandler := interceptor.GetMetricsHandlerFromContext(ctx, logger).WithTags(metrics.OperationTag(metrics.HistoryGetRawHistoryScope))
-			metrics.ServiceErrIncompleteHistoryCounter.With(metricsHandler).Record(1)
-			logger.Error("getHistory error",
-				tag.WorkflowNamespaceID(namespaceID.String()),
-				tag.WorkflowID(execution.GetWorkflowId()),
-				tag.WorkflowRunID(execution.GetRunId()),
-				tag.Error(err))
-			return nil, nil, err
-		}
-
-		if len(transientWorkflowTaskInfo.HistorySuffix) > 0 {
-			blob, err := shardContext.GetPayloadSerializer().SerializeEvents(transientWorkflowTaskInfo.HistorySuffix)
-			if err != nil {
-				return nil, nil, err
+		// Check if we should include transient/speculative events
+		if shouldIncludeTransientOrSpeculativeTasks(ctx, transientWorkflowTaskInfo) {
+			if err := ValidateTransientWorkflowTaskEvents(nextEventID, transientWorkflowTaskInfo); err != nil {
+				logger := shardContext.GetLogger()
+				metricsHandler := interceptor.GetMetricsHandlerFromContext(ctx, logger).WithTags(metrics.OperationTag(metrics.HistoryGetRawHistoryScope))
+				metrics.ServiceErrIncompleteHistoryCounter.With(metricsHandler).Record(1)
+			} else {
+				if len(transientWorkflowTaskInfo.HistorySuffix) > 0 {
+					blob, err := shardContext.GetPayloadSerializer().SerializeEvents(transientWorkflowTaskInfo.HistorySuffix)
+					if err != nil {
+						return nil, nil, err
+					}
+					rawHistory = append(rawHistory, blob)
+				}
 			}
-			rawHistory = append(rawHistory, blob)
 		}
 	}
 
@@ -235,16 +232,19 @@ func GetHistory(
 			tag.Error(err))
 	}
 	if len(nextPageToken) == 0 && transientWorkflowTaskInfo != nil {
-		if err := validateTransientWorkflowTaskEvents(nextEventID, transientWorkflowTaskInfo); err != nil {
-			metrics.ServiceErrIncompleteHistoryCounter.With(metricsHandler).Record(1)
-			logger.Error("getHistory error",
-				tag.WorkflowNamespaceID(namespaceID.String()),
-				tag.WorkflowID(execution.GetWorkflowId()),
-				tag.WorkflowRunID(execution.GetRunId()),
-				tag.Error(err))
+		// Check if we should include transient/speculative events
+		if shouldIncludeTransientOrSpeculativeTasks(ctx, transientWorkflowTaskInfo) {
+			if err := ValidateTransientWorkflowTaskEvents(nextEventID, transientWorkflowTaskInfo); err != nil {
+				metrics.ServiceErrIncompleteHistoryCounter.With(metricsHandler).Record(1)
+				// Don't append events, but don't fail request
+			} else {
+				if len(transientWorkflowTaskInfo.HistorySuffix) > 0 {
+					// Validation passed, append events
+					// Append the transient workflow task events once we are done enumerating everything from the events table
+					historyEvents = append(historyEvents, transientWorkflowTaskInfo.HistorySuffix...)
+				}
+			}
 		}
-		// Append the transient workflow task events once we are done enumerating everything from the events table
-		historyEvents = append(historyEvents, transientWorkflowTaskInfo.HistorySuffix...)
 	}
 
 	if err := ProcessOutgoingSearchAttributes(
@@ -386,7 +386,58 @@ func ProcessOutgoingSearchAttributes(
 	return nil
 }
 
-func validateTransientWorkflowTaskEvents(
+// shouldIncludeTransientOrSpeculativeTasks determines if transient/speculative events should be included.
+// This function is called only when on the last page of history pagination (nextToken is empty).
+func shouldIncludeTransientOrSpeculativeTasks(
+	ctx context.Context,
+	tranOrSpecEvents *historyspb.TransientWorkflowTaskInfo,
+) bool {
+	return len(tranOrSpecEvents.GetHistorySuffix()) > 0 &&
+		clientSupportsTranOrSpecEvents(ctx) &&
+		areValidTransientOrSpecEvents(tranOrSpecEvents)
+}
+
+func areValidTransientOrSpecEvents(tranOrSpecEvents *historyspb.TransientWorkflowTaskInfo) bool {
+	events := tranOrSpecEvents.GetHistorySuffix()
+	if len(events) == 0 || len(events) > 2 {
+		return false
+	}
+
+	// First must be WFT_SCHEDULED
+	if events[0].GetEventType() != enumspb.EVENT_TYPE_WORKFLOW_TASK_SCHEDULED {
+		return false
+	}
+
+	// If 2 events, second must be WFT_STARTED immediately after
+	if len(events) == 2 {
+		if events[1].GetEventType() != enumspb.EVENT_TYPE_WORKFLOW_TASK_STARTED {
+			return false
+		}
+		if events[1].GetEventId() != events[0].GetEventId()+1 {
+			return false
+		}
+	}
+
+	return true
+}
+
+// clientSupportsTranOrSpecEvents detects if client supports transient events
+// Default to include transient events for clients, only CLI and UI are
+// explicitly excluded for backward compatability
+func clientSupportsTranOrSpecEvents(ctx context.Context) bool {
+	clientName, _ := headers.GetClientNameAndVersion(ctx)
+
+	switch clientName {
+	case headers.ClientNameCLI, headers.ClientNameUI:
+		return false
+	default:
+		return true
+	}
+}
+
+// ValidateTransientWorkflowTaskEvents validates that transient workflow task events have sequential event IDs
+// starting from the given offset. Returns an error if any event ID doesn't match the expected sequence.
+func ValidateTransientWorkflowTaskEvents(
 	eventIDOffset int64,
 	transientWorkflowTaskInfo *historyspb.TransientWorkflowTaskInfo,
 ) error {