Refactors to avoid calling the activity CHASM component directly

Author: davidporter-id-au

chasm/lib/activity/activity.go

@@ -71,7 +71,9 @@ type ActivityStore interface {
 	// (true for START_TO_CLOSE and HEARTBEAT timeouts; false for SCHEDULE_TO_START/CLOSE).
 	OnActivityTimedOut(ctx chasm.MutableContext, act *Activity, timeoutFailure *failurepb.Failure, needsStartedEvent bool) error
 	// OnActivityCanceled is called when an activity reaches the Canceled terminal state.
-	OnActivityCanceled(ctx chasm.MutableContext, act *Activity) error
+	// needsStartedEvent indicates whether an ActivityTaskStarted event must also be written
+	// (true when the activity was started before being canceled; false for not-yet-started activities).
+	OnActivityCanceled(ctx chasm.MutableContext, act *Activity, needsStartedEvent bool) error
 	// OnActivityTerminated is called when an activity reaches the Terminated terminal state.
 	OnActivityTerminated(ctx chasm.MutableContext, act *Activity) error
 }
@@ -328,7 +330,7 @@ func (a *Activity) OnActivityTimedOut(ctx chasm.MutableContext, _ *Activity, _ *
 }
 
 // OnActivityCanceled implements ActivityStore for standalone activities.
-func (a *Activity) OnActivityCanceled(ctx chasm.MutableContext, _ *Activity) error {
+func (a *Activity) OnActivityCanceled(ctx chasm.MutableContext, _ *Activity, _ bool) error {
 	return callback.ScheduleStandbyCallbacks(ctx, a.Callbacks)
 }
 

chasm/lib/activity/statemachine.go

@@ -243,6 +243,11 @@ var TransitionFailed = chasm.NewTransition(
 			return err
 		}
 
+		outcome := a.Outcome.Get(ctx)
+		outcome.Variant = &activitypb.ActivityOutcome_Failed_{
+			Failed: &activitypb.ActivityOutcome_Failed{Failure: req.GetFailure()},
+		}
+
 		a.emitOnFailedMetrics(ctx, event.metricsHandler)
 
 		return a.StoreOrSelf(ctx).OnActivityFailed(ctx, a)
@@ -339,7 +344,8 @@ var TransitionCanceled = chasm.NewTransition(
 
 		a.emitOnCanceledMetrics(ctx, event.handler, event.fromStatus)
 
-		return a.StoreOrSelf(ctx).OnActivityCanceled(ctx, a)
+		needsStartedEvent := event.fromStatus != activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED
+		return a.StoreOrSelf(ctx).OnActivityCanceled(ctx, a, needsStartedEvent)
 	},
 )
 

chasm/lib/activity/statemachine_test.go

@@ -562,7 +562,7 @@ func TestTransitionFailed(t *testing.T) {
 	require.NotNil(t, heartbeatState.GetRecordedTime())
 	protorequire.ProtoEqual(t, failure, attemptState.GetLastFailureDetails().GetFailure())
 	require.NotNil(t, attemptState.GetLastFailureDetails().GetTime())
-	require.Nil(t, outcome.GetFailed())
+	protorequire.ProtoEqual(t, failure, outcome.GetFailed().GetFailure())
 }
 
 func TestTransitionTerminated(t *testing.T) {

chasm/lib/workflow/workflow.go

@@ -9,11 +9,13 @@ import (
 	historypb "go.temporal.io/api/history/v1"
 	"go.temporal.io/api/serviceerror"
 	workflowpb "go.temporal.io/api/workflow/v1"
+	"go.temporal.io/server/api/historyservice/v1"
 	"go.temporal.io/server/chasm"
 	"go.temporal.io/server/chasm/lib/activity"
 	"go.temporal.io/server/chasm/lib/callback"
 	callbackspb "go.temporal.io/server/chasm/lib/callback/gen/callbackpb/v1"
 	"go.temporal.io/server/chasm/lib/nexusoperation"
+	"go.temporal.io/server/common/metrics"
 	"go.temporal.io/server/service/history/historybuilder"
 	"google.golang.org/protobuf/types/known/emptypb"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -176,11 +178,46 @@ func (w *Workflow) OnActivityTimedOut(ctx chasm.MutableContext, act *activity.Ac
 }
 
 // OnActivityCanceled implements ActivityStore for workflow-embedded activities.
-// No ActivityTaskCanceled history event is written on the forward path in this prototype.
-func (w *Workflow) OnActivityCanceled(ctx chasm.MutableContext, act *activity.Activity) error {
-	activityID := act.GetActivityId()
-	delete(w.Activities, activityID)
-	return w.ScheduleWorkflowTask()
+// When needsStartedEvent is true (activity was started before being canceled), writes
+// ActivityTaskStarted + ActivityTaskCanceled history events; Apply() handles cleanup.
+// When needsStartedEvent is false (activity was never started), the caller is responsible
+// for writing the ActivityTaskCanceled event (via AddActivityTaskCanceledEventCHASM).
+func (w *Workflow) OnActivityCanceled(ctx chasm.MutableContext, act *activity.Activity, needsStartedEvent bool) error {
+	if !needsStartedEvent {
+		// Not-started case: ActivityTaskCanceled is written by AddActivityTaskCanceledEventCHASM.
+		// Just clean up the activity and schedule a WFT.
+		activityID := act.GetActivityId()
+		delete(w.Activities, activityID)
+		return w.ScheduleWorkflowTask()
+	}
+
+	scheduledEventID := act.GetScheduledEventId()
+	attempt := act.LastAttempt.Get(ctx)
+	startedEvent, err := addAndApplyHistoryEvent[ActivityTaskStartedEventDefinition](w, ctx, func(e *historypb.HistoryEvent) {
+		e.Attributes = &historypb.HistoryEvent_ActivityTaskStartedEventAttributes{
+			ActivityTaskStartedEventAttributes: &historypb.ActivityTaskStartedEventAttributes{
+				ScheduledEventId: scheduledEventID,
+				Attempt:          attempt.GetCount(),
+				RequestId:        attempt.GetStartRequestId(),
+				Identity:         attempt.GetLastWorkerIdentity(),
+			},
+		}
+	})
+	if err != nil {
+		return err
+	}
+	cancelDetails := act.Outcome.Get(ctx).GetFailed().GetFailure().GetCanceledFailureInfo().GetDetails()
+	_, err = addAndApplyHistoryEvent[ActivityTaskCanceledEventDefinition](w, ctx, func(e *historypb.HistoryEvent) {
+		e.Attributes = &historypb.HistoryEvent_ActivityTaskCanceledEventAttributes{
+			ActivityTaskCanceledEventAttributes: &historypb.ActivityTaskCanceledEventAttributes{
+				ScheduledEventId: scheduledEventID,
+				StartedEventId:   startedEvent.GetEventId(),
+				Details:          cancelDetails,
+				Identity:         attempt.GetLastWorkerIdentity(),
+			},
+		}
+	})
+	return err
 }
 
 // OnActivityTerminated implements ActivityStore for workflow-embedded activities.
@@ -243,6 +280,54 @@ func (w *Workflow) AddCompletionCallbacks(
 	return nil
 }
 
+// RegisterScheduledActivity applies TransitionScheduled to act and adds it to the workflow's
+// Activities map. This is the single entry point for scheduling embedded activities so that
+// the transition trigger stays inside the workflow package rather than in the caller.
+func (w *Workflow) RegisterScheduledActivity(
+	ctx chasm.MutableContext,
+	activityID string,
+	act *activity.Activity,
+) error {
+	if err := activity.TransitionScheduled.Apply(act, ctx, nil); err != nil {
+		return err
+	}
+	w.AddEmbeddedActivity(ctx, activityID, act)
+	return nil
+}
+
+// RequestCancelEmbeddedActivity finds the embedded activity with scheduledEventID, applies the
+// cancel transition, and returns (wasNotStarted, found, error). found=false means no CHASM
+// activity has that scheduledEventID; the caller should fall back to the legacy path.
+func (w *Workflow) RequestCancelEmbeddedActivity(
+	ctx chasm.MutableContext,
+	scheduledEventID int64,
+	identity string,
+) (wasNotStarted bool, found bool, err error) {
+	act := w.FindActivityByScheduledEventID(ctx, scheduledEventID)
+	if act == nil {
+		return false, false, nil
+	}
+	wasNotStarted, err = act.RequestCancelFromWorkflowTask(ctx, identity)
+	return wasNotStarted, true, err
+}
+
+// CompleteEmbeddedActivityByID finds the activity by activityID, applies TransitionCompleted,
+// and returns (true, nil) on success. Returns (false, nil) if no CHASM activity with that ID
+// exists; the caller should fall back to the legacy path.
+func (w *Workflow) CompleteEmbeddedActivityByID(
+	ctx chasm.MutableContext,
+	activityID string,
+	req *historyservice.RespondActivityTaskCompletedRequest,
+	metricsHandler metrics.Handler,
+) (bool, error) {
+	actField, ok := w.Activities[activityID]
+	if !ok {
+		return false, nil
+	}
+	act := actField.Get(ctx)
+	return true, activity.TransitionCompleted.Apply(act, ctx, activity.NewCompleteEvent(req, metricsHandler))
+}
+
 // AddEmbeddedActivity adds a CHASM activity sub-component to the workflow, keyed by SDK-provided activity ID.
 func (w *Workflow) AddEmbeddedActivity(
 	ctx chasm.MutableContext,

service/history/workflow/mutable_state_impl.go

@@ -4175,16 +4175,12 @@ func (ms *MutableStateImpl) AddActivityTaskScheduledEventCHASM(
 		}),
 	}
 
-	// 4. Call TransitionScheduled first (generates ActivityDispatchTask and timeout pure tasks),
-	// following the same pattern as nexus operations where the transition is applied before
-	// the component is added to the parent's map field.
-	if err := chasmactivity.TransitionScheduled.Apply(act, chasmCtx, nil); err != nil {
+	// 4. Register: applies TransitionScheduled (generates ActivityDispatchTask and timeout pure
+	// tasks) then adds to the workflow's Activities map, keyed by activity ID.
+	if err := wf.RegisterScheduledActivity(chasmCtx, command.GetActivityId(), act); err != nil {
 		return 0, err
 	}
 
-	// 5. Add to the workflow's Activities map — this registers it in the CHASM tree, keyed by activity ID.
-	wf.AddEmbeddedActivity(chasmCtx, command.GetActivityId(), act)
-
 	// Accumulate the size of this CHASM activity so GetApproximatePersistedSize includes it
 	// before CloseTransaction serializes it into approximateSize.
 	ms.chasmPendingSize += command.GetInput().Size() + command.GetHeader().Size() + act.Size()
@@ -4652,8 +4648,11 @@ func (ms *MutableStateImpl) AddActivityTaskCancelRequestedEventCHASM(
 		return nil, false, err
 	}
 
-	act := wf.FindActivityByScheduledEventID(chasmCtx, scheduledEventID)
-	if act == nil {
+	wasNotStarted, found, err := wf.RequestCancelEmbeddedActivity(chasmCtx, scheduledEventID, identity)
+	if err != nil {
+		return nil, false, err
+	}
+	if !found {
 		// Return ErrCHASMActivityNotFound so the caller can fall back to the legacy
 		// AddActivityTaskCancelRequestedEvent path (backwards-compat guard for executions
 		// that scheduled the activity before EnableCHASMActivityPrototype was enabled).
@@ -4662,11 +4661,6 @@ func (ms *MutableStateImpl) AddActivityTaskCancelRequestedEventCHASM(
 
 	actCancelReqEvent := ms.hBuilder.AddActivityTaskCancelRequestedEvent(workflowTaskCompletedEventID, scheduledEventID)
 
-	wasNotStarted, err := act.RequestCancelFromWorkflowTask(chasmCtx, identity)
-	if err != nil {
-		return nil, false, err
-	}
-
 	// Caller is responsible for writing ActivityTaskCanceled when wasNotStarted is true,
 	// following the same pattern as the legacy AddActivityTaskCancelRequestedEvent path.
 	return actCancelReqEvent, wasNotStarted, nil
@@ -4716,21 +4710,15 @@ func (ms *MutableStateImpl) RespondCHASMActivityCompletedByID(activityID string,
 	if err != nil {
 		return false, err
 	}
-	actField, ok := wf.Activities[activityID]
-	if !ok {
-		return false, nil
-	}
-	act := actField.Get(chasmCtx)
 	metricsHandler := ms.metricsHandler.WithTags(metrics.OperationTag(metrics.HistoryRespondActivityTaskCompletedScope))
-	// Build a minimal RespondActivityTaskCompletedRequest for the transition.
 	req := &historyservice.RespondActivityTaskCompletedRequest{
 		NamespaceId: ms.GetNamespaceEntry().ID().String(),
 		CompleteRequest: &workflowservice.RespondActivityTaskCompletedRequest{
 			Result:   result,
 			Identity: identity,
 		},
 	}
-	return true, chasmactivity.TransitionCompleted.Apply(act, chasmCtx, chasmactivity.NewCompleteEvent(req, metricsHandler))
+	return wf.CompleteEmbeddedActivityByID(chasmCtx, activityID, req, metricsHandler)
 }
 
 func (ms *MutableStateImpl) AddWorkerCommandsTasks(commands []*workerpb.WorkerCommand, controlQueue string) error {

tests/chasm_activity_test.go

@@ -364,7 +364,7 @@ func (s *ChasmActivitySuite) TestChasmActivity_PendingActivityDescribe() {
 
 // assertContainsEventType is a helper that asserts the given event type appears at least once
 // in the history events slice.
-func (s *ChasmActivitySuite) assertContainsEventType(events []*historypb.HistoryEvent, eventType enumspb.EventType, msgAndArgs ...interface{}) {
+func (s *ChasmActivitySuite) assertContainsEventType(events []*historypb.HistoryEvent, eventType enumspb.EventType, msgAndArgs ...any) {
 	s.T().Helper()
 	for _, e := range events {
 		if e.EventType == eventType {