Added the channel machines and notification delivery to the external Core.

crates/protos/protos/local/temporal/sdk/core/workflow_commands/workflow_commands.proto

@@ -55,9 +55,31 @@ message WorkflowCommand {
         ExternalStreamFinalized external_stream_finalized = 26;
         WorkflowOutputStreamCommit workflow_output_stream_commit = 27;
         WorkflowOutputStreamBuffered workflow_output_stream_buffered = 28;
+        SubscribeNotificationChannel subscribe_notification_channel = 29;
+        UnsubscribeNotificationChannel unsubscribe_notification_channel = 30;
     }
 }
 
+// End this workflow's subscription to a notification channel. Notifications
+// already recorded on a scheduled event still reach that Workflow Task.
+//
+// The server records the event whether or not the run held a subscription, so
+// every command has an event to match on replay.
+message UnsubscribeNotificationChannel {
+    // Name of the channel, scoped to the namespace.
+    string channel = 1;
+}
+
+// Subscribe this workflow to a notification channel, so the scheduled event of
+// each later Workflow Task carries the notifications folded for it.
+//
+// The notifications live in History rather than arriving by a side channel, so
+// a replay reads the same ones the live run saw.
+message SubscribeNotificationChannel {
+    // Name of the channel, scoped to the namespace.
+    string channel = 1;
+}
+
 // Commits an observation delta for external streams.
 //
 // Emitted on *every* completion path where replay-visible stream state changed -- which includes

crates/protos/src/protos/mod.rs

@@ -1805,6 +1805,18 @@ pub mod coresdk {
                 }
             }
 
+            impl Display for SubscribeNotificationChannel {
+                fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+                    write!(f, "SubscribeNotificationChannel({})", self.channel)
+                }
+            }
+
+            impl Display for UnsubscribeNotificationChannel {
+                fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+                    write!(f, "UnsubscribeNotificationChannel({})", self.channel)
+                }
+            }
+
             impl Display for WorkflowOutputStreamBuffered {
                 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
                     write!(
@@ -2056,6 +2068,26 @@ pub mod temporal {
                         }
                     }
 
+                    impl From<workflow_commands::SubscribeNotificationChannel> for Attributes {
+                        fn from(s: workflow_commands::SubscribeNotificationChannel) -> Self {
+                            Self::SubscribeNotificationChannelCommandAttributes(
+                                SubscribeNotificationChannelCommandAttributes {
+                                    channel: s.channel,
+                                },
+                            )
+                        }
+                    }
+
+                    impl From<workflow_commands::UnsubscribeNotificationChannel> for Attributes {
+                        fn from(u: workflow_commands::UnsubscribeNotificationChannel) -> Self {
+                            Self::UnsubscribeNotificationChannelCommandAttributes(
+                                UnsubscribeNotificationChannelCommandAttributes {
+                                    channel: u.channel,
+                                },
+                            )
+                        }
+                    }
+
                     impl From<workflow_commands::StartTimer> for command::Attributes {
                         fn from(s: workflow_commands::StartTimer) -> Self {
                             Self::StartTimerCommandAttributes(StartTimerCommandAttributes {

crates/sdk-core/src/core_tests/external_streams.rs

@@ -51,6 +51,7 @@ use temporalio_common::{
             command::v1::command,
             common::v1::Payload,
             enums::v1::{CommandType, EventType},
+            notification::v1::Notification,
             query::v1::WorkflowQuery,
             workflowservice::v1::RespondWorkflowTaskCompletedResponse,
         },
@@ -3297,6 +3298,169 @@ async fn an_ordinary_signal_still_reaches_its_handler() {
     worker.drain_pollers_and_shutdown().await;
 }
 
+// --- notifications on the scheduled event ------------------------------------
+
+fn channel_notification(channel: &str, counter: i64) -> Notification {
+    Notification {
+        channel: channel.to_string(),
+        position: format!("{counter}-0").into_bytes(),
+        counter,
+        metadata: HashMap::new(),
+        ..Default::default()
+    }
+}
+
+/// The notifications of every `NotificationsReceived` job in the activation, one entry per job.
+fn notification_jobs(activation: &WorkflowActivation) -> Vec<Vec<Notification>> {
+    activation
+        .jobs
+        .iter()
+        .filter_map(|j| match &j.variant {
+            Some(workflow_activation_job::Variant::NotificationsReceived(n)) => {
+                Some(n.notifications.clone())
+            }
+            _ => None,
+        })
+        .collect()
+}
+
+/// A history whose second task fails after being scheduled with `n1` and is retried with `n2`.
+fn failed_then_retried_history(n1: Vec<Notification>, n2: Vec<Notification>) -> TestHistoryBuilder {
+    let mut t = TestHistoryBuilder::default();
+    t.add_wfe_started_with_wft_timeout(Duration::from_secs(300));
+    t.add_full_wf_task();
+    t.add_workflow_task_scheduled_with_notifications(n1);
+    t.add_workflow_task_started();
+    t.add_workflow_task_failed_with_failure(
+        temporalio_common::protos::temporal::api::enums::v1::WorkflowTaskFailedCause::Unspecified,
+        Default::default(),
+    );
+    t.add_workflow_task_scheduled_with_notifications(n2);
+    t.add_workflow_task_started();
+    t
+}
+
+#[rstest::rstest]
+#[case::live(false)]
+#[case::replay(true)]
+#[tokio::test]
+async fn a_failed_task_hands_its_notifications_to_the_retry(#[case] cold: bool) {
+    // The server clears a listener's pending notifications when it puts them on a scheduled
+    // event, so a retry's event carries only what arrived after the failed attempt. The failed
+    // attempt's event stays in History, and its notifications reach lang on the retry, folded
+    // with the retry's own into one job: one per channel, the highest counter kept.
+    let n1 = vec![channel_notification("tokens", 41)];
+    let n2 = vec![
+        channel_notification("tokens", 42),
+        channel_notification("events", 3),
+    ];
+    let t = failed_then_retried_history(n1, n2.clone());
+    let batches = if cold {
+        let resp = hist_to_poll_resp(&t, "fakeid".to_owned(), ResponseType::AllHistory);
+        vec![ResponseType::Raw(resp.resp)]
+    } else {
+        vec![ResponseType::ToTaskNum(1), ResponseType::ToTaskNum(2)]
+    };
+    let mut mock = build_mock_pollers(MockPollCfg::from_resp_batches(
+        "fakeid",
+        t,
+        batches,
+        mock_worker_client(),
+    ));
+    mock.worker_cfg(|w| {
+        w.task_types = WorkerTaskTypes::workflow_only();
+        w.max_cached_workflows = 1;
+    });
+    let worker = mock_worker(mock);
+    let first = worker.poll_workflow_activation().await.unwrap();
+    assert_eq!(first.is_replaying, cold);
+    assert_eq!(
+        notification_jobs(&first),
+        Vec::<Vec<Notification>>::new(),
+        "the notifications belong to the retried task, got {:?}",
+        first.jobs
+    );
+    worker
+        .complete_workflow_activation(WorkflowActivationCompletion::empty(first.run_id))
+        .await
+        .unwrap();
+
+    let retried = worker.poll_workflow_activation().await.unwrap();
+    assert!(!retried.is_replaying, "got jobs {:?}", retried.jobs);
+    assert_eq!(notification_jobs(&retried), vec![n2]);
+    worker
+        .complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
+            retried.run_id,
+            vec![CompleteWorkflowExecution::default().into()],
+        ))
+        .await
+        .unwrap();
+    worker.drain_pollers_and_shutdown().await;
+}
+
+#[tokio::test]
+async fn replaying_a_notified_scheduled_event_yields_the_same_job() {
+    // History is the record, so a cold worker hands lang the same job in the same activation the
+    // live run saw it in. The first task issued no command, which is the shape Core would
+    // otherwise treat as a heartbeat and fold into the next task, moving the job one activation
+    // early. The Signal keeps the replayed task apart from the live one.
+    let notifications = vec![channel_notification("tokens", 42)];
+    let mut t = TestHistoryBuilder::default();
+    t.add_wfe_started_with_wft_timeout(Duration::from_secs(300));
+    t.add_full_wf_task();
+    t.add_workflow_task_scheduled_with_notifications(notifications.clone());
+    t.add_workflow_task_started();
+    t.add_workflow_task_completed();
+    t.add_we_signaled("a-user-signal", vec![]);
+    t.add_workflow_task_scheduled_and_started();
+
+    let cold = hist_to_poll_resp(&t, "fakeid".to_owned(), ResponseType::AllHistory);
+    let mut mock = build_mock_pollers(MockPollCfg::from_resp_batches(
+        "fakeid",
+        t,
+        [ResponseType::Raw(cold.resp)],
+        mock_worker_client(),
+    ));
+    mock.worker_cfg(|w| {
+        w.task_types = WorkerTaskTypes::workflow_only();
+        w.max_cached_workflows = 1;
+    });
+    let worker = mock_worker(mock);
+
+    let first = worker.poll_workflow_activation().await.unwrap();
+    assert!(first.is_replaying);
+    assert_eq!(
+        notification_jobs(&first),
+        Vec::<Vec<Notification>>::new(),
+        "the notifications belong to the second task, got {:?}",
+        first.jobs
+    );
+    worker
+        .complete_workflow_activation(WorkflowActivationCompletion::empty(first.run_id))
+        .await
+        .unwrap();
+
+    let replayed = worker.poll_workflow_activation().await.unwrap();
+    assert!(replayed.is_replaying, "got jobs {:?}", replayed.jobs);
+    assert_eq!(notification_jobs(&replayed), vec![notifications]);
+    worker
+        .complete_workflow_activation(WorkflowActivationCompletion::empty(replayed.run_id))
+        .await
+        .unwrap();
+
+    let live = worker.poll_workflow_activation().await.unwrap();
+    assert!(!live.is_replaying, "got jobs {:?}", live.jobs);
+    assert_eq!(notification_jobs(&live), Vec::<Vec<Notification>>::new());
+    worker
+        .complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
+            live.run_id,
+            vec![CompleteWorkflowExecution::default().into()],
+        ))
+        .await
+        .unwrap();
+    worker.drain_pollers_and_shutdown().await;
+}
+
 // --- marker emission (C9, C14b) ----------------------------------------------
 
 /// A history for a Workflow Task that records a marker and then starts a timer.

crates/sdk-core/src/replay/history_builder.rs

@@ -114,6 +114,43 @@ impl TestHistoryBuilder {
         self.workflow_task_scheduled_event_id = self.add_by_type(EventType::WorkflowTaskScheduled);
     }
 
+    /// Add a workflow task scheduled event carrying channel notifications.
+    pub fn add_workflow_task_scheduled_with_notifications(
+        &mut self,
+        notifications: Vec<
+            temporalio_common::protos::temporal::api::notification::v1::Notification,
+        >,
+    ) {
+        self.workflow_task_scheduled_event_id = self.add(WorkflowTaskScheduledEventAttributes {
+            notifications,
+            ..Default::default()
+        });
+    }
+
+    /// Add the event a subscribe-notification-channel command produces.
+    pub fn add_notification_channel_subscribed(&mut self, channel: &str) -> i64 {
+        let attrs = WorkflowNotificationChannelSubscribedEventAttributes {
+            workflow_task_completed_event_id: self.previous_task_completed_id,
+            channel: channel.to_string(),
+        };
+        self.add(attrs)
+    }
+
+    /// Add the event an unsubscribe-notification-channel command produces. A zero
+    /// `subscribed_event_id` is what the server records when the run held no subscription.
+    pub fn add_notification_channel_unsubscribed(
+        &mut self,
+        channel: &str,
+        subscribed_event_id: i64,
+    ) -> i64 {
+        let attrs = WorkflowNotificationChannelUnsubscribedEventAttributes {
+            workflow_task_completed_event_id: self.previous_task_completed_id,
+            channel: channel.to_string(),
+            subscribed_event_id,
+        };
+        self.add(attrs)
+    }
+
     /// Add a workflow task started event.
     pub fn add_workflow_task_started(&mut self) {
         self.final_workflow_task_started_event_id = self.add(WorkflowTaskStartedEventAttributes {

crates/sdk-core/src/worker/workflow/history_update.rs

@@ -741,10 +741,14 @@ fn find_end_index_of_next_wft_seq(
                     if let Some(next_next_event) = events.get(ix + 2) {
                         if !saw_command
                             && next_next_event.event_type() == EventType::WorkflowTaskScheduled
+                            && !scheduled_with_notifications(next_next_event)
                         {
                             // If we've never seen an interesting event and the next two events are
                             // a completion followed immediately again by scheduled, then this is a
-                            // WFT heartbeat and also doesn't conclude the sequence.
+                            // WFT heartbeat and also doesn't conclude the sequence. A scheduled
+                            // event with notifications is work lang saw in its own activation
+                            // live, so folding it into this one on replay would hand lang the
+                            // notifications one activation early.
                             continue;
                         } else {
                             // If we see an update accepted command after WFT completed, we want to
@@ -801,6 +805,15 @@ fn find_end_index_of_next_wft_seq(
     NextWFTSeqEndIndex::Incomplete(last_index)
 }
 
+/// A task scheduled with notifications was handed them as an activation of its own, so it
+/// cannot be folded into the heartbeat chain before it the way an empty task is.
+fn scheduled_with_notifications(e: &HistoryEvent) -> bool {
+    matches!(
+        &e.attributes,
+        Some(Attributes::WorkflowTaskScheduledEventAttributes(a)) if !a.notifications.is_empty()
+    )
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/sdk-core/src/worker/workflow/machines/mod.rs

@@ -16,13 +16,19 @@ mod modify_workflow_properties_state_machine;
 mod nexus_operation_state_machine;
 mod patch_state_machine;
 mod signal_external_state_machine;
+mod subscribe_notification_channel_state_machine;
 mod timer_state_machine;
+mod unsubscribe_notification_channel_state_machine;
 mod update_state_machine;
 mod upsert_search_attributes_state_machine;
 mod workflow_task_state_machine;
 
+#[cfg(test)]
+mod subscribe_notification_channel_state_machine_tests;
 #[cfg(test)]
 mod transition_coverage;
+#[cfg(test)]
+mod unsubscribe_notification_channel_state_machine_tests;
 
 pub(crate) use temporalio_common::fsm_trait::MachineError;
 pub(crate) use workflow_machines::{MachinesWFTResponseContent, WorkflowMachines};
@@ -48,6 +54,7 @@ use std::{
     convert::{TryFrom, TryInto},
     fmt::{Debug, Display},
 };
+use subscribe_notification_channel_state_machine::SubscribeNotificationChannelMachine;
 use temporalio_common::{
     fsm_trait::{StateMachine, TransitionResult},
     protos::temporal::api::{
@@ -57,6 +64,7 @@ use temporalio_common::{
 };
 use temporalio_macros::fsm;
 use timer_state_machine::TimerMachine;
+use unsubscribe_notification_channel_state_machine::UnsubscribeNotificationChannelMachine;
 use update_state_machine::UpdateMachine;
 use upsert_search_attributes_state_machine::UpsertSearchAttributesMachine;
 use workflow_machines::MachineResponse;
@@ -85,6 +93,8 @@ enum Machines {
     ModifyWorkflowPropertiesMachine,
     UpdateMachine,
     NexusOperationMachine,
+    SubscribeNotificationChannelMachine,
+    UnsubscribeNotificationChannelMachine,
 }
 
 /// Extends [rustfsm::StateMachine] with some functionality specific to the temporal SDK.