@temporalio/core-bridge 1.6.0 → 1.7.1

package/sdk-core/core/src/core_tests/determinism.rs

@@ -1,4 +1,5 @@
 use crate::{
+    internal_flags::CoreInternalFlags,
     replay::DEFAULT_WORKFLOW_TYPE,
     test_help::{canned_histories, mock_sdk, mock_sdk_cfg, MockPollCfg, ResponseType},
     worker::client::mocks::mock_workflow_client,
@@ -8,8 +9,16 @@ use std::{
     time::Duration,
 };
 use temporal_client::WorkflowOptions;
-use temporal_sdk::{WfContext, WorkflowResult};
-use temporal_sdk_core_protos::temporal::api::enums::v1::WorkflowTaskFailedCause;
+use temporal_sdk::{
+    ActivityOptions, ChildWorkflowOptions, LocalActivityOptions, WfContext, WorkflowResult,
+};
+use temporal_sdk_core_protos::{
+    temporal::api::{
+        enums::v1::{EventType, WorkflowTaskFailedCause},
+        failure::v1::Failure,
+    },
+    TestHistoryBuilder, DEFAULT_ACTIVITY_TYPE,
+};
 
 static DID_FAIL: AtomicBool = AtomicBool::new(false);
 pub async fn timer_wf_fails_once(ctx: WfContext) -> WorkflowResult<()> {
@@ -105,3 +114,204 @@ async fn test_wf_task_rejected_properly_due_to_nondeterminism(#[case] use_cache:
     // timer and proceed without restarting
     assert_eq!(2, started_count.load(Ordering::Relaxed));
 }
+
+#[rstest::rstest]
+#[tokio::test]
+async fn activity_id_or_type_change_is_nondeterministic(
+    #[values(true, false)] use_cache: bool,
+    #[values(true, false)] id_change: bool,
+    #[values(true, false)] local_act: bool,
+) {
+    let wf_id = "fakeid";
+    let wf_type = DEFAULT_WORKFLOW_TYPE;
+    let mut t = if local_act {
+        canned_histories::single_local_activity("1")
+    } else {
+        canned_histories::single_activity("1")
+    };
+    t.set_flags_first_wft(&[CoreInternalFlags::IdAndTypeDeterminismChecks as u32], &[]);
+    let mock = mock_workflow_client();
+    let mut mh = MockPollCfg::from_resp_batches(
+        wf_id,
+        t,
+        // Two polls are needed, since the first will fail
+        [ResponseType::AllHistory, ResponseType::AllHistory],
+        mock,
+    );
+    // We should see one wft failure which has nondeterminism cause
+    mh.num_expected_fails = 1;
+    mh.expect_fail_wft_matcher = Box::new(move |_, cause, f| {
+        let should_contain = if id_change {
+            "does not match activity id"
+        } else {
+            "does not match activity type"
+        };
+        matches!(cause, WorkflowTaskFailedCause::NonDeterministicError)
+            && matches!(f, Some(Failure {
+                message,
+                ..
+            }) if message.contains(should_contain))
+    });
+    let mut worker = mock_sdk_cfg(mh, |cfg| {
+        if use_cache {
+            cfg.max_cached_workflows = 2;
+        }
+    });
+
+    worker.register_wf(wf_type.to_owned(), move |ctx: WfContext| async move {
+        if local_act {
+            ctx.local_activity(if id_change {
+                LocalActivityOptions {
+                    activity_id: Some("I'm bad and wrong!".to_string()),
+                    activity_type: DEFAULT_ACTIVITY_TYPE.to_string(),
+                    ..Default::default()
+                }
+            } else {
+                LocalActivityOptions {
+                    activity_type: "not the default act type".to_string(),
+                    ..Default::default()
+                }
+            })
+            .await;
+        } else {
+            ctx.activity(if id_change {
+                ActivityOptions {
+                    activity_id: Some("I'm bad and wrong!".to_string()),
+                    activity_type: DEFAULT_ACTIVITY_TYPE.to_string(),
+                    ..Default::default()
+                }
+            } else {
+                ActivityOptions {
+                    activity_type: "not the default act type".to_string(),
+                    ..Default::default()
+                }
+            })
+            .await;
+        }
+        Ok(().into())
+    });
+
+    worker
+        .submit_wf(
+            wf_id.to_owned(),
+            wf_type.to_owned(),
+            vec![],
+            WorkflowOptions::default(),
+        )
+        .await
+        .unwrap();
+    worker.run_until_done().await.unwrap();
+}
+
+#[rstest::rstest]
+#[tokio::test]
+async fn child_wf_id_or_type_change_is_nondeterministic(
+    #[values(true, false)] use_cache: bool,
+    #[values(true, false)] id_change: bool,
+) {
+    let wf_id = "fakeid";
+    let wf_type = DEFAULT_WORKFLOW_TYPE;
+    let mut t = canned_histories::single_child_workflow("1");
+    t.set_flags_first_wft(&[CoreInternalFlags::IdAndTypeDeterminismChecks as u32], &[]);
+    let mock = mock_workflow_client();
+    let mut mh = MockPollCfg::from_resp_batches(
+        wf_id,
+        t,
+        // Two polls are needed, since the first will fail
+        [ResponseType::AllHistory, ResponseType::AllHistory],
+        mock,
+    );
+    // We should see one wft failure which has nondeterminism cause
+    mh.num_expected_fails = 1;
+    mh.expect_fail_wft_matcher = Box::new(move |_, cause, f| {
+        let should_contain = if id_change {
+            "does not match child workflow id"
+        } else {
+            "does not match child workflow type"
+        };
+        matches!(cause, WorkflowTaskFailedCause::NonDeterministicError)
+            && matches!(f, Some(Failure {
+                message,
+                ..
+            }) if message.contains(should_contain))
+    });
+    let mut worker = mock_sdk_cfg(mh, |cfg| {
+        if use_cache {
+            cfg.max_cached_workflows = 2;
+        }
+    });
+
+    worker.register_wf(wf_type.to_owned(), move |ctx: WfContext| async move {
+        ctx.child_workflow(if id_change {
+            ChildWorkflowOptions {
+                workflow_id: "I'm bad and wrong!".to_string(),
+                workflow_type: DEFAULT_ACTIVITY_TYPE.to_string(),
+                ..Default::default()
+            }
+        } else {
+            ChildWorkflowOptions {
+                workflow_id: "1".to_string(),
+                workflow_type: "not the child wf type".to_string(),
+                ..Default::default()
+            }
+        })
+        .start(&ctx)
+        .await;
+        Ok(().into())
+    });
+
+    worker
+        .submit_wf(
+            wf_id.to_owned(),
+            wf_type.to_owned(),
+            vec![],
+            WorkflowOptions::default(),
+        )
+        .await
+        .unwrap();
+    worker.run_until_done().await.unwrap();
+}
+
+/// Repros a situation where if, upon completing a task there is some internal error which causes
+/// us to want to auto-fail the workflow task while there is also an outstanding eviction, the wf
+/// would get evicted but then try to send some info down the completion channel afterward, causing
+/// a panic.
+#[tokio::test]
+async fn repro_channel_missing_because_nondeterminism() {
+    for _ in 1..50 {
+        let wf_id = "fakeid";
+        let wf_type = DEFAULT_WORKFLOW_TYPE;
+        let mut t = TestHistoryBuilder::default();
+        t.add_by_type(EventType::WorkflowExecutionStarted);
+        t.add_full_wf_task();
+        t.add_has_change_marker("patch-1", false);
+        let _ts = t.add_by_type(EventType::TimerStarted);
+        t.add_workflow_task_scheduled_and_started();
+
+        let mock = mock_workflow_client();
+        let mut mh =
+            MockPollCfg::from_resp_batches(wf_id, t, [1.into(), ResponseType::AllHistory], mock);
+        mh.num_expected_fails = 1;
+        let mut worker = mock_sdk_cfg(mh, |cfg| {
+            cfg.max_cached_workflows = 2;
+            cfg.ignore_evicts_on_shutdown = false;
+        });
+
+        worker.register_wf(wf_type.to_owned(), move |ctx: WfContext| async move {
+            ctx.patched("wrongid");
+            ctx.timer(Duration::from_secs(1)).await;
+            Ok(().into())
+        });
+
+        worker
+            .submit_wf(
+                wf_id.to_owned(),
+                wf_type.to_owned(),
+                vec![],
+                WorkflowOptions::default(),
+            )
+            .await
+            .unwrap();
+        worker.run_until_done().await.unwrap();
+    }
+}

package/sdk-core/core/src/core_tests/local_activities.rs

@@ -8,6 +8,7 @@ use crate::{
     worker::{client::mocks::mock_workflow_client, LEGACY_QUERY_ID},
 };
 use anyhow::anyhow;
+use crossbeam::queue::SegQueue;
 use futures::{future::join_all, FutureExt};
 use std::{
     collections::HashMap,
@@ -22,26 +23,32 @@ use temporal_client::WorkflowOptions;
 use temporal_sdk::{
     ActContext, ActivityCancelledError, LocalActivityOptions, WfContext, WorkflowResult,
 };
-use temporal_sdk_core_api::Worker;
+use temporal_sdk_core_api::{
+    errors::{PollActivityError, PollWfError},
+    Worker,
+};
 use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::ActivityExecutionResult,
         workflow_activation::{workflow_activation_job, WorkflowActivationJob},
-        workflow_commands::{ActivityCancellationType, QueryResult, QuerySuccess},
+        workflow_commands::{
+            ActivityCancellationType, QueryResult, QuerySuccess, ScheduleLocalActivity,
+        },
         workflow_completion::WorkflowActivationCompletion,
         ActivityTaskCompletion, AsJsonPayloadExt,
     },
     temporal::api::{
         common::v1::RetryPolicy,
         enums::v1::{EventType, TimeoutType, WorkflowTaskFailedCause},
-        failure::v1::Failure,
+        failure::v1::{failure::FailureInfo, Failure},
         query::v1::WorkflowQuery,
     },
+    DEFAULT_ACTIVITY_TYPE,
 };
 use temporal_sdk_core_test_utils::{
     schedule_local_activity_cmd, start_timer_cmd, WorkerTestHelpers,
 };
-use tokio::sync::Barrier;
+use tokio::{join, sync::Barrier};
 
 async fn echo(_ctx: ActContext, e: String) -> anyhow::Result<String> {
     Ok(e)
@@ -59,7 +66,7 @@ async fn local_act_two_wfts_before_marker(#[case] replay: bool, #[case] cached:
     let mut t = TestHistoryBuilder::default();
     t.add_by_type(EventType::WorkflowExecutionStarted);
     t.add_full_wf_task();
-    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
     t.add_full_wf_task();
     t.add_local_activity_result_marker(1, "1", b"echo".into());
     t.add_timer_fired(timer_started_event_id, "1".to_string());
@@ -84,7 +91,7 @@ async fn local_act_two_wfts_before_marker(#[case] replay: bool, #[case] cached:
         DEFAULT_WORKFLOW_TYPE.to_owned(),
         |ctx: WfContext| async move {
             let la = ctx.local_activity(LocalActivityOptions {
-                activity_type: "echo".to_string(),
+                activity_type: DEFAULT_ACTIVITY_TYPE.to_string(),
                 input: "hi".as_json_payload().expect("serializes fine"),
                 ..Default::default()
             });
@@ -93,7 +100,7 @@ async fn local_act_two_wfts_before_marker(#[case] replay: bool, #[case] cached:
             Ok(().into())
         },
     );
-    worker.register_activity("echo", echo);
+    worker.register_activity(DEFAULT_ACTIVITY_TYPE, echo);
     worker
         .submit_wf(
             wf_id.to_owned(),
@@ -125,11 +132,10 @@ pub async fn local_act_fanout_wf(ctx: WfContext) -> WorkflowResult<()> {
 
 #[tokio::test]
 async fn local_act_many_concurrent() {
-    crate::telemetry::test_telem_console();
     let mut t = TestHistoryBuilder::default();
     t.add_by_type(EventType::WorkflowExecutionStarted);
     t.add_full_wf_task();
-    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
     t.add_full_wf_task();
     for i in 1..=50 {
         t.add_local_activity_result_marker(i, &i.to_string(), b"echo".into());
@@ -301,7 +307,7 @@ async fn local_act_retry_long_backoff_uses_timer() {
         "1",
         Failure::application_failure("la failed".to_string(), false),
     );
-    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
     t.add_timer_fired(timer_started_event_id, "1".to_string());
     t.add_full_wf_task();
     t.add_local_activity_fail_marker(
@@ -309,7 +315,7 @@ async fn local_act_retry_long_backoff_uses_timer() {
         "2",
         Failure::application_failure("la failed".to_string(), false),
     );
-    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
     t.add_timer_fired(timer_started_event_id, "2".to_string());
     t.add_full_wf_task();
     t.add_workflow_execution_completed();
@@ -329,7 +335,7 @@ async fn local_act_retry_long_backoff_uses_timer() {
         |ctx: WfContext| async move {
             let la_res = ctx
                 .local_activity(LocalActivityOptions {
-                    activity_type: "echo".to_string(),
+                    activity_type: DEFAULT_ACTIVITY_TYPE.to_string(),
                     input: "hi".as_json_payload().expect("serializes fine"),
                     retry_policy: RetryPolicy {
                         initial_interval: Some(prost_dur!(from_millis(65))),
@@ -348,9 +354,12 @@ async fn local_act_retry_long_backoff_uses_timer() {
             Ok(().into())
         },
     );
-    worker.register_activity("echo", move |_ctx: ActContext, _: String| async move {
-        Result::<(), _>::Err(anyhow!("Oh no I failed!"))
-    });
+    worker.register_activity(
+        DEFAULT_ACTIVITY_TYPE,
+        move |_ctx: ActContext, _: String| async move {
+            Result::<(), _>::Err(anyhow!("Oh no I failed!"))
+        },
+    );
     worker
         .submit_wf(
             wf_id.to_owned(),
@@ -410,7 +419,7 @@ async fn local_act_command_immediately_follows_la_marker() {
     t.add_full_wf_task();
     t.add_full_wf_task();
     t.add_local_activity_result_marker(1, "1", "done".into());
-    t.add_get_event_id(EventType::TimerStarted, None);
+    t.add_by_type(EventType::TimerStarted);
     t.add_full_wf_task();
 
     let wf_id = "fakeid";
@@ -498,7 +507,7 @@ async fn query_during_wft_heartbeat_doesnt_accidentally_fail_to_continue_heartbe
             task.run_id,
             schedule_local_activity_cmd(
                 1,
-                "act-id",
+                "1",
                 ActivityCancellationType::TryCancel,
                 Duration::from_secs(60),
             ),
@@ -564,15 +573,11 @@ async fn la_resolve_during_legacy_query_does_not_combine(#[case] impossible_quer
     // never happen, but there was an issue where an LA resolving could trigger that.
     let wfid = "fake_wf_id";
     let mut t = TestHistoryBuilder::default();
-    let wes_short_wft_timeout = default_wes_attribs();
-    t.add(
-        EventType::WorkflowExecutionStarted,
-        wes_short_wft_timeout.into(),
-    );
+    t.add(default_wes_attribs());
     // Since we don't send queries with start workflow, need one workflow task of something else
     // b/c we want to get an activation with a job and a nonlegacy query
     t.add_full_wf_task();
-    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
     t.add_timer_fired(timer_started_event_id, "1".to_string());
 
     // nonlegacy query got here & LA started here
@@ -760,6 +765,15 @@ async fn test_schedule_to_start_timeout() {
                 })
                 .await;
             assert_eq!(la_res.timed_out(), Some(TimeoutType::ScheduleToStart));
+            let rfail = la_res.unwrap_failure();
+            assert_matches!(
+                rfail.failure_info,
+                Some(FailureInfo::ActivityFailureInfo(_))
+            );
+            assert_matches!(
+                rfail.cause.unwrap().failure_info,
+                Some(FailureInfo::TimeoutFailureInfo(_))
+            );
             Ok(().into())
         },
     );
@@ -808,7 +822,7 @@ async fn test_schedule_to_start_timeout_not_based_on_original_time(
             deets.backoff = Some(prost_dur!(from_secs(100)));
         },
     );
-    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
     t.add_timer_fired(timer_started_event_id, "1".to_string());
     t.add_workflow_task_scheduled_and_started();
 
@@ -874,7 +888,7 @@ async fn wft_failure_cancels_running_las() {
     let mut t = TestHistoryBuilder::default();
     t.add_wfe_started_with_wft_timeout(Duration::from_millis(200));
     t.add_full_wf_task();
-    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
     t.add_timer_fired(timer_started_event_id, "1".to_string());
     t.add_workflow_task_scheduled_and_started();
 
@@ -888,7 +902,7 @@ async fn wft_failure_cancels_running_las() {
         DEFAULT_WORKFLOW_TYPE.to_owned(),
         |ctx: WfContext| async move {
             let la_handle = ctx.local_activity(LocalActivityOptions {
-                activity_type: "echo".to_string(),
+                activity_type: DEFAULT_ACTIVITY_TYPE.to_string(),
                 input: "hi".as_json_payload().expect("serializes fine"),
                 ..Default::default()
             });
@@ -902,13 +916,16 @@ async fn wft_failure_cancels_running_las() {
             Ok(().into())
         },
     );
-    worker.register_activity("echo", move |ctx: ActContext, _: String| async move {
-        let res = tokio::time::timeout(Duration::from_millis(500), ctx.cancelled()).await;
-        if res.is_err() {
-            panic!("Activity must be cancelled!!!!");
-        }
-        Result::<(), _>::Err(ActivityCancelledError::default().into())
-    });
+    worker.register_activity(
+        DEFAULT_ACTIVITY_TYPE,
+        move |ctx: ActContext, _: String| async move {
+            let res = tokio::time::timeout(Duration::from_millis(500), ctx.cancelled()).await;
+            if res.is_err() {
+                panic!("Activity must be cancelled!!!!");
+            }
+            Result::<(), _>::Err(ActivityCancelledError::default().into())
+        },
+    );
     worker
         .submit_wf(
             wf_id.to_owned(),
@@ -944,9 +961,7 @@ async fn resolved_las_not_recorded_if_wft_fails_many_times() {
         mock,
     );
     mh.num_expected_fails = 2;
-    mh.completion_asserts = Some(Box::new(|_| {
-        panic!("should never successfully complete a WFT");
-    }));
+    mh.num_expected_completions = Some(0.into());
     let mut worker = mock_sdk_cfg(mh, |w| w.max_cached_workflows = 1);
 
     worker.register_wf(
@@ -976,3 +991,135 @@ async fn resolved_las_not_recorded_if_wft_fails_many_times() {
         .unwrap();
     worker.run_until_done().await.unwrap();
 }
+
+#[tokio::test]
+async fn local_act_records_nonfirst_attempts_ok() {
+    let mut t = TestHistoryBuilder::default();
+    let wft_timeout = Duration::from_millis(200);
+    t.add_wfe_started_with_wft_timeout(wft_timeout);
+    t.add_full_wf_task();
+    t.add_full_wf_task();
+    t.add_full_wf_task();
+    t.add_workflow_task_scheduled_and_started();
+
+    let wf_id = "fakeid";
+    let mock = mock_workflow_client();
+    let mut mh = MockPollCfg::from_resp_batches(wf_id, t, [1, 2, 3], mock);
+    let nonfirst_counts = Arc::new(SegQueue::new());
+    let nfc_c = nonfirst_counts.clone();
+    mh.completion_asserts = Some(Box::new(move |c| {
+        nfc_c.push(
+            c.metering_metadata
+                .nonfirst_local_activity_execution_attempts,
+        );
+    }));
+    let mut worker = mock_sdk_cfg(mh, |wc| {
+        wc.max_cached_workflows = 1;
+        wc.max_outstanding_workflow_tasks = 1;
+    });
+
+    worker.register_wf(
+        DEFAULT_WORKFLOW_TYPE.to_owned(),
+        |ctx: WfContext| async move {
+            ctx.local_activity(LocalActivityOptions {
+                activity_type: "echo".to_string(),
+                input: "hi".as_json_payload().expect("serializes fine"),
+                retry_policy: RetryPolicy {
+                    initial_interval: Some(prost_dur!(from_millis(10))),
+                    backoff_coefficient: 1.0,
+                    maximum_interval: None,
+                    maximum_attempts: 0,
+                    non_retryable_error_types: vec![],
+                },
+                ..Default::default()
+            })
+            .await;
+            Ok(().into())
+        },
+    );
+    worker.register_activity("echo", move |_ctx: ActContext, _: String| async move {
+        Result::<(), _>::Err(anyhow!("I fail"))
+    });
+    worker
+        .submit_wf(
+            wf_id.to_owned(),
+            DEFAULT_WORKFLOW_TYPE.to_owned(),
+            vec![],
+            WorkflowOptions::default(),
+        )
+        .await
+        .unwrap();
+    worker.run_until_done().await.unwrap();
+    // 3 workflow tasks
+    assert_eq!(nonfirst_counts.len(), 3);
+    // First task's non-first count should, of course, be 0
+    assert_eq!(nonfirst_counts.pop().unwrap(), 0);
+    // Next two, some nonzero amount which could vary based on test load
+    assert!(nonfirst_counts.pop().unwrap() > 0);
+    assert!(nonfirst_counts.pop().unwrap() > 0);
+}
+
+#[tokio::test]
+async fn local_activities_can_be_delivered_during_shutdown() {
+    let wfid = "fake_wf_id";
+    let mut t = TestHistoryBuilder::default();
+    t.add_wfe_started_with_wft_timeout(Duration::from_millis(200));
+    t.add_full_wf_task();
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
+    t.add_timer_fired(timer_started_event_id, "1".to_string());
+    t.add_workflow_task_scheduled_and_started();
+
+    let mock = mock_workflow_client();
+    let mut mock = single_hist_mock_sg(
+        wfid,
+        t,
+        [ResponseType::ToTaskNum(1), ResponseType::AllHistory],
+        mock,
+        true,
+    );
+    mock.worker_cfg(|wc| wc.max_cached_workflows = 1);
+    let core = mock_worker(mock);
+
+    let task = core.poll_workflow_activation().await.unwrap();
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+        task.run_id,
+        start_timer_cmd(1, Duration::from_secs(1)),
+    ))
+    .await
+    .unwrap();
+
+    let task = core.poll_workflow_activation().await.unwrap();
+    // Initiate shutdown once we have the WF activation, but before replying that we want to do an
+    // LA
+    core.initiate_shutdown();
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+        task.run_id,
+        ScheduleLocalActivity {
+            seq: 1,
+            activity_id: "1".to_string(),
+            activity_type: "test_act".to_string(),
+            start_to_close_timeout: Some(prost_dur!(from_secs(30))),
+            ..Default::default()
+        }
+        .into(),
+    ))
+    .await
+    .unwrap();
+
+    let wf_poller = async { core.poll_workflow_activation().await };
+
+    let at_poller = async {
+        let act_task = core.poll_activity_task().await.unwrap();
+        core.complete_activity_task(ActivityTaskCompletion {
+            task_token: act_task.task_token,
+            result: Some(ActivityExecutionResult::ok(vec![1].into())),
+        })
+        .await
+        .unwrap();
+        core.poll_activity_task().await
+    };
+
+    let (wf_r, act_r) = join!(wf_poller, at_poller);
+    assert_matches!(wf_r.unwrap_err(), PollWfError::ShutDown);
+    assert_matches!(act_r.unwrap_err(), PollActivityError::ShutDown);
+}

package/sdk-core/core/src/core_tests/queries.rs

@@ -23,7 +23,7 @@ use temporal_sdk_core_protos::{
     },
     temporal::api::{
         common::v1::Payload,
-        enums::v1::EventType,
+        enums::v1::{CommandType, EventType},
         failure::v1::Failure,
         history::v1::{history_event, ActivityTaskCancelRequestedEventAttributes, History},
         query::v1::WorkflowQuery,
@@ -148,10 +148,11 @@ async fn legacy_query(#[case] include_history: bool) {
 }
 
 #[rstest::rstest]
-#[case::one_query(1)]
-#[case::multiple_queries(3)]
 #[tokio::test]
-async fn new_queries(#[case] num_queries: usize) {
+async fn new_queries(
+    #[values(1, 3)] num_queries: usize,
+    #[values(false, true)] query_results_after_complete: bool,
+) {
     let wfid = "fake_wf_id";
     let query_resp = "response";
     let t = canned_histories::single_timer("1");
@@ -174,7 +175,18 @@ async fn new_queries(#[case] num_queries: usize) {
     }]);
     let mut mock_client = mock_workflow_client();
     mock_client.expect_respond_legacy_query().times(0);
-    let mut mock = single_hist_mock_sg(wfid, t, tasks, mock_client, true);
+    let mut mh = MockPollCfg::from_resp_batches(wfid, t, tasks, mock_workflow_client());
+    mh.completion_asserts = Some(Box::new(move |c| {
+        // If the completion is the one ending the workflow, make sure it includes the query resps
+        if c.commands[0].command_type() == CommandType::CompleteWorkflowExecution {
+            assert_eq!(c.query_responses.len(), num_queries);
+        } else if c.commands[0].command_type() == CommandType::StartTimer {
+            // first reply, no queries here.
+        } else {
+            panic!("Unexpected command in response")
+        }
+    }));
+    let mut mock = build_mock_pollers(mh);
     mock.worker_cfg(|wc| wc.max_cached_workflows = 10);
     let core = mock_worker(mock);
 
@@ -203,8 +215,13 @@ async fn new_queries(#[case] num_queries: usize) {
             }
         );
     }
-    let mut qresults: Vec<_> = (1..=num_queries)
-        .map(|i| {
+
+    let mut commands = vec![];
+    if query_results_after_complete {
+        commands.push(CompleteWorkflowExecution { result: None }.into());
+    }
+    for i in 1..=num_queries {
+        commands.push(
             QueryResult {
                 query_id: format!("q{i}"),
                 variant: Some(
@@ -214,13 +231,15 @@ async fn new_queries(#[case] num_queries: usize) {
                     .into(),
                 ),
             }
-            .into()
-        })
-        .collect();
-    qresults.push(CompleteWorkflowExecution { result: None }.into());
+            .into(),
+        );
+    }
+    if !query_results_after_complete {
+        commands.push(CompleteWorkflowExecution { result: None }.into());
+    }
     core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
         task.run_id,
-        qresults,
+        commands,
     ))
     .await
     .unwrap();
@@ -769,11 +788,10 @@ async fn legacy_query_combined_with_timer_fire_repro() {
     t.add_by_type(EventType::WorkflowExecutionStarted);
     t.add_full_wf_task();
     let scheduled_event_id = t.add_activity_task_scheduled("1");
-    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
     t.add_timer_fired(timer_started_event_id, "1".to_string());
     t.add_full_wf_task();
     t.add(
-        EventType::ActivityTaskCancelRequested,
         history_event::Attributes::ActivityTaskCancelRequestedEventAttributes(
             ActivityTaskCancelRequestedEventAttributes {
                 scheduled_event_id,