@temporalio/core-bridge 1.11.2 → 1.11.4

package/sdk-core/core/src/abstractions.rs

@@ -11,7 +11,8 @@ use std::{
     },
 };
 use temporal_sdk_core_api::worker::{
-    SlotKind, SlotReservationContext, SlotSupplier, SlotSupplierPermit,
+    SlotKind, SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext, SlotSupplier,
+    SlotSupplierPermit, WorkflowSlotKind,
 };
 use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
@@ -35,6 +36,18 @@ pub(crate) struct MeteredPermitDealer<SK: SlotKind> {
     /// there will need to be some associated refactoring.
     max_permits: Option<usize>,
     metrics_ctx: MetricsContext,
+    /// Only applies to permit dealers for workflow tasks. True if this permit dealer is associated
+    /// with a sticky queue poller.
+    is_sticky_poller: bool,
+    context_data: Arc<PermitDealerContextData>,
+}
+
+#[derive(Clone, Debug)]
+#[cfg_attr(test, derive(Default))]
+pub(crate) struct PermitDealerContextData {
+    pub(crate) task_queue: String,
+    pub(crate) worker_identity: String,
+    pub(crate) worker_build_id: String,
 }
 
 impl<SK> MeteredPermitDealer<SK>
@@ -45,6 +58,7 @@ where
         supplier: Arc<dyn SlotSupplier<SlotKind = SK> + Send + Sync>,
         metrics_ctx: MetricsContext,
         max_permits: Option<usize>,
+        context_data: Arc<PermitDealerContextData>,
     ) -> Self {
         Self {
             supplier,
@@ -52,6 +66,8 @@ where
             extant_permits: watch::channel(0),
             metrics_ctx,
             max_permits,
+            is_sticky_poller: false,
+            context_data,
         }
     }
 
@@ -91,6 +107,10 @@ where
         }
     }
 
+    pub(crate) fn get_extant_count_rcv(&self) -> watch::Receiver<usize> {
+        self.extant_permits.1.clone()
+    }
+
     fn build_owned(&self, res: SlotSupplierPermit) -> OwnedMeteredSemPermit<SK> {
         self.unused_claimants.fetch_add(1, Ordering::Release);
         self.extant_permits.0.send_modify(|ep| *ep += 1);
@@ -117,14 +137,17 @@ where
         mrc(false);
 
         OwnedMeteredSemPermit {
-            _inner: res,
             unused_claimants: Some(self.unused_claimants.clone()),
+            release_ctx: ReleaseCtx {
+                permit: res,
+                stored_info: None,
+            },
             use_fn: Box::new(move |info| {
                 supp_c.mark_slot_used(info);
                 metric_rec(false)
             }),
-            release_fn: Box::new(move || {
-                supp_c_c.release_slot();
+            release_fn: Box::new(move |info| {
+                supp_c_c.release_slot(info);
                 ep_tx_c.send_modify(|ep| *ep -= 1);
                 mrc(true)
             }),
@@ -132,10 +155,67 @@ where
     }
 }
 
+impl MeteredPermitDealer<WorkflowSlotKind> {
+    pub(crate) fn into_sticky(mut self) -> Self {
+        self.is_sticky_poller = true;
+        self
+    }
+}
+
 impl<SK: SlotKind> SlotReservationContext for MeteredPermitDealer<SK> {
+    fn task_queue(&self) -> &str {
+        &self.context_data.task_queue
+    }
+
+    fn worker_identity(&self) -> &str {
+        &self.context_data.worker_identity
+    }
+
+    fn worker_build_id(&self) -> &str {
+        &self.context_data.worker_build_id
+    }
+
     fn num_issued_slots(&self) -> usize {
         *self.extant_permits.1.borrow()
     }
+
+    fn is_sticky(&self) -> bool {
+        self.is_sticky_poller
+    }
+}
+
+struct UseCtx<'a, SK: SlotKind> {
+    stored_info: &'a SK::Info,
+    permit: &'a SlotSupplierPermit,
+}
+
+impl<'a, SK: SlotKind> SlotMarkUsedContext for UseCtx<'a, SK> {
+    type SlotKind = SK;
+
+    fn permit(&self) -> &SlotSupplierPermit {
+        self.permit
+    }
+
+    fn info(&self) -> &<Self::SlotKind as SlotKind>::Info {
+        self.stored_info
+    }
+}
+
+struct ReleaseCtx<SK: SlotKind> {
+    permit: SlotSupplierPermit,
+    stored_info: Option<SK::Info>,
+}
+
+impl<SK: SlotKind> SlotReleaseContext for ReleaseCtx<SK> {
+    type SlotKind = SK;
+
+    fn permit(&self) -> &SlotSupplierPermit {
+        &self.permit
+    }
+
+    fn info(&self) -> Option<&<Self::SlotKind as SlotKind>::Info> {
+        self.stored_info.as_ref()
+    }
 }
 
 /// A version of [MeteredPermitDealer] that can be closed and supports waiting for close to complete.
@@ -240,20 +320,22 @@ impl<SK: SlotKind> Drop for TrackedOwnedMeteredSemPermit<SK> {
 /// Wraps an [SlotSupplierPermit] to update metrics & when it's dropped
 #[clippy::has_significant_drop]
 pub(crate) struct OwnedMeteredSemPermit<SK: SlotKind> {
-    _inner: SlotSupplierPermit,
     /// See [MeteredPermitDealer::unused_claimants]. If present when dropping, used to decrement the
     /// count.
     unused_claimants: Option<Arc<AtomicUsize>>,
+    /// The actual [SlotSupplierPermit] is stored in here
+    release_ctx: ReleaseCtx<SK>,
+    #[allow(clippy::type_complexity)] // not really tho, bud
+    use_fn: Box<dyn Fn(&UseCtx<SK>) + Send + Sync>,
     #[allow(clippy::type_complexity)] // not really tho, bud
-    use_fn: Box<dyn Fn(SK::Info<'_>) + Send + Sync>,
-    release_fn: Box<dyn Fn() + Send + Sync>,
+    release_fn: Box<dyn Fn(&ReleaseCtx<SK>) + Send + Sync>,
 }
 impl<SK: SlotKind> Drop for OwnedMeteredSemPermit<SK> {
     fn drop(&mut self) {
         if let Some(uc) = self.unused_claimants.take() {
             uc.fetch_sub(1, Ordering::Release);
         }
-        (self.release_fn)()
+        (self.release_fn)(&self.release_ctx);
     }
 }
 impl<SK: SlotKind> Debug for OwnedMeteredSemPermit<SK> {
@@ -264,11 +346,16 @@ impl<SK: SlotKind> Debug for OwnedMeteredSemPermit<SK> {
 impl<SK: SlotKind> OwnedMeteredSemPermit<SK> {
     /// Should be called once this permit is actually being "used" for the work it was meant to
     /// permit.
-    pub(crate) fn into_used(mut self, info: SK::Info<'_>) -> UsedMeteredSemPermit<SK> {
+    pub(crate) fn into_used(mut self, info: SK::Info) -> UsedMeteredSemPermit<SK> {
         if let Some(uc) = self.unused_claimants.take() {
             uc.fetch_sub(1, Ordering::Release);
         }
-        (self.use_fn)(info);
+        let ctx = UseCtx {
+            stored_info: &info,
+            permit: &self.release_ctx.permit,
+        };
+        (self.use_fn)(&ctx);
+        self.release_ctx.stored_info = Some(info);
         UsedMeteredSemPermit(self)
     }
 }
@@ -298,6 +385,7 @@ pub(crate) mod tests {
             Arc::new(FixedSizeSlotSupplier::new(size)),
             MetricsContext::no_op(),
             None,
+            Arc::new(Default::default()),
         )
     }
 

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

@@ -9,7 +9,7 @@ use crate::{
     worker::client::mocks::{mock_manual_workflow_client, mock_workflow_client},
     ActivityHeartbeat, Worker,
 };
-use futures::FutureExt;
+use futures_util::FutureExt;
 use itertools::Itertools;
 use std::{
     cell::RefCell,
@@ -145,6 +145,7 @@ async fn heartbeats_report_cancels_only_once() {
         .returning(|_, _| {
             Ok(RecordActivityTaskHeartbeatResponse {
                 cancel_requested: true,
+                activity_paused: false,
             })
         });
     mock_client
@@ -270,6 +271,7 @@ async fn activity_cancel_interrupts_poll() {
             async {
                 Ok(RecordActivityTaskHeartbeatResponse {
                     cancel_requested: true,
+                    activity_paused: false,
                 })
             }
             .boxed()
@@ -390,10 +392,12 @@ async fn many_concurrent_heartbeat_cancels() {
                 if calls < 5 {
                     Ok(RecordActivityTaskHeartbeatResponse {
                         cancel_requested: false,
+                        activity_paused: false,
                     })
                 } else {
                     Ok(RecordActivityTaskHeartbeatResponse {
                         cancel_requested: true,
+                        activity_paused: false,
                     })
                 }
             }
@@ -512,6 +516,7 @@ async fn can_heartbeat_acts_during_shutdown() {
         .returning(|_, _| {
             Ok(RecordActivityTaskHeartbeatResponse {
                 cancel_requested: false,
+                activity_paused: false,
             })
         });
     mock_client
@@ -565,6 +570,7 @@ async fn complete_act_with_fail_flushes_heartbeat() {
             *lsp.borrow_mut() = payload;
             Ok(RecordActivityTaskHeartbeatResponse {
                 cancel_requested: false,
+                activity_paused: false,
             })
         });
     mock_client
@@ -674,7 +680,7 @@ async fn no_eager_activities_requested_when_worker_options_disable_remote_activi
     let mut mock_poller = mock_manual_poller();
     mock_poller
         .expect_poll()
-        .returning(|| futures::future::pending().boxed());
+        .returning(|| futures_util::future::pending().boxed());
     mock.set_act_poller(Box::new(mock_poller));
     mock.worker_cfg(|wc| {
         wc.max_cached_workflows = 2;

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

@@ -9,7 +9,7 @@ use crate::{
 };
 use anyhow::anyhow;
 use crossbeam_queue::SegQueue;
-use futures::{future::join_all, FutureExt};
+use futures_util::{future::join_all, FutureExt};
 use std::{
     collections::HashMap,
     ops::Sub,

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

@@ -15,8 +15,8 @@ use crate::{
     worker::client::mocks::{mock_manual_workflow_client, mock_workflow_client},
     Worker,
 };
-use futures::FutureExt;
-use once_cell::sync::Lazy;
+use futures_util::FutureExt;
+use std::sync::LazyLock;
 use std::time::Duration;
 use temporal_sdk_core_api::Worker as WorkerTrait;
 use temporal_sdk_core_protos::coresdk::workflow_completion::WorkflowActivationCompletion;
@@ -46,7 +46,7 @@ async fn after_shutdown_server_is_not_polled() {
 }
 
 // Better than cloning a billion arcs...
-static BARR: Lazy<Barrier> = Lazy::new(|| Barrier::new(3));
+static BARR: LazyLock<Barrier> = LazyLock::new(|| Barrier::new(3));
 
 #[tokio::test]
 async fn shutdown_interrupts_both_polls() {

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

@@ -1,4 +1,5 @@
 use crate::{
+    prost_dur,
     test_help::{
         build_mock_pollers, hist_to_poll_resp, mock_worker, MockPollCfg, PollWFTRespExt,
         ResponseType,
@@ -10,7 +11,8 @@ use temporal_sdk_core_protos::{
     coresdk::{
         workflow_activation::{workflow_activation_job, WorkflowActivationJob},
         workflow_commands::{
-            update_response::Response, CompleteWorkflowExecution, ScheduleActivity, UpdateResponse,
+            update_response::Response, CompleteWorkflowExecution, ScheduleActivity, StartTimer,
+            UpdateResponse,
         },
         workflow_completion::WorkflowActivationCompletion,
     },
@@ -41,17 +43,25 @@ async fn replay_with_empty_first_task() {
     mock.worker_cfg(|wc| wc.max_cached_workflows = 1);
     let core = mock_worker(mock);
 
+    // In this task imagine we are waiting on the first update being sent, hence no commands come
+    // out, and on replay the first activation should only be init.
     let task = core.poll_workflow_activation().await.unwrap();
     assert_matches!(
         task.jobs.as_slice(),
-        [
-            WorkflowActivationJob {
-                variant: Some(workflow_activation_job::Variant::InitializeWorkflow(_)),
-            },
-            WorkflowActivationJob {
-                variant: Some(workflow_activation_job::Variant::DoUpdate(_)),
-            },
-        ]
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::InitializeWorkflow(_)),
+        },]
+    );
+    core.complete_workflow_activation(WorkflowActivationCompletion::empty(task.run_id))
+        .await
+        .unwrap();
+
+    let task = core.poll_workflow_activation().await.unwrap();
+    assert_matches!(
+        task.jobs.as_slice(),
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::DoUpdate(_)),
+        },]
     );
     core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
         task.run_id,
@@ -229,3 +239,88 @@ async fn speculative_wft_with_command_event() {
     );
     core.complete_execution(&task.run_id).await;
 }
+
+#[tokio::test]
+async fn replay_with_signal_and_update_same_task() {
+    // Imitating a signal creating a command before update validator runs
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+    t.add_we_signaled("hi", vec![]);
+    t.add_full_wf_task();
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
+    let accept_id = t.add_update_accepted("upd1", "update");
+    t.add_timer_fired(timer_started_event_id, "1".to_string());
+    t.add_full_wf_task();
+    t.add_update_completed(accept_id);
+    t.add_workflow_execution_completed();
+
+    let mock = MockPollCfg::from_resps(t, [ResponseType::AllHistory]);
+    let mut mock = build_mock_pollers(mock);
+    mock.worker_cfg(|wc| wc.max_cached_workflows = 1);
+    let core = mock_worker(mock);
+
+    // In this task imagine we are waiting on the first update being sent, hence no commands come
+    // out, and on replay the first activation should only be init.
+    let task = core.poll_workflow_activation().await.unwrap();
+    assert_matches!(
+        task.jobs.as_slice(),
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::InitializeWorkflow(_)),
+        },]
+    );
+    core.complete_workflow_activation(WorkflowActivationCompletion::empty(task.run_id))
+        .await
+        .unwrap();
+
+    let task = core.poll_workflow_activation().await.unwrap();
+    assert_matches!(
+        task.jobs.as_slice(),
+        [
+            WorkflowActivationJob {
+                variant: Some(workflow_activation_job::Variant::SignalWorkflow(_)),
+            },
+            WorkflowActivationJob {
+                variant: Some(workflow_activation_job::Variant::DoUpdate(_)),
+            }
+        ]
+    );
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
+        task.run_id,
+        vec![
+            StartTimer {
+                seq: 1,
+                start_to_fire_timeout: Some(prost_dur!(from_secs(1))),
+            }
+            .into(),
+            UpdateResponse {
+                protocol_instance_id: "upd1".to_string(),
+                response: Some(Response::Accepted(())),
+            }
+            .into(),
+        ],
+    ))
+    .await
+    .unwrap();
+
+    let task = core.poll_workflow_activation().await.unwrap();
+    assert_matches!(
+        task.jobs.as_slice(),
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::FireTimer(_)),
+        },]
+    );
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
+        task.run_id,
+        vec![
+            UpdateResponse {
+                protocol_instance_id: "upd1".to_string(),
+                response: Some(Response::Completed(Payload::default())),
+            }
+            .into(),
+            CompleteWorkflowExecution { result: None }.into(),
+        ],
+    ))
+    .await
+    .unwrap();
+}

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

@@ -4,8 +4,13 @@ use crate::{
         build_fake_worker, build_mock_pollers, canned_histories, mock_worker, test_worker_cfg,
         MockPollCfg, MockWorkerInputs, MocksHolder, ResponseType, WorkerExt,
     },
-    worker,
-    worker::client::mocks::mock_workflow_client,
+    worker::{
+        self,
+        client::{
+            mocks::{mock_workflow_client, DEFAULT_TEST_CAPABILITIES, DEFAULT_WORKERS_REGISTRY},
+            MockWorkerClient,
+        },
+    },
     PollActivityError, PollWfError,
 };
 use futures_util::{stream, stream::StreamExt};
@@ -18,7 +23,7 @@ use temporal_sdk_core_protos::{
         workflow_completion::WorkflowActivationCompletion,
     },
     temporal::api::workflowservice::v1::{
-        PollWorkflowTaskQueueResponse, RespondWorkflowTaskCompletedResponse,
+        PollWorkflowTaskQueueResponse, RespondWorkflowTaskCompletedResponse, ShutdownWorkerResponse,
     },
 };
 use temporal_sdk_core_test_utils::start_timer_cmd;
@@ -295,3 +300,67 @@ async fn worker_can_shutdown_after_never_polling_ok(#[values(true, false)] poll_
         break;
     }
 }
+
+#[rstest::rstest]
+#[case::ok(true, true)]
+#[case::best_effort(true, false)]
+#[case::not_sticky(false, true)]
+#[tokio::test]
+async fn worker_shutdown_api(#[case] use_cache: bool, #[case] api_success: bool) {
+    // Manually need to create MockWorkerClient because we want to specify
+    // the expected number of calls for shutdown_worker.
+    // This will no longer be needed if
+    // https://github.com/asomers/mockall/issues/283 is implemented.
+    let mut mock = MockWorkerClient::new();
+    mock.expect_capabilities()
+        .returning(|| Some(*DEFAULT_TEST_CAPABILITIES));
+    mock.expect_workers()
+        .returning(|| DEFAULT_WORKERS_REGISTRY.clone());
+    mock.expect_is_mock().returning(|| true);
+    if use_cache {
+        if api_success {
+            mock.expect_shutdown_worker()
+                .times(1)
+                .returning(|_| Ok(ShutdownWorkerResponse {}));
+        } else {
+            // worker.shutdown() should succeed even if shutdown_worker fails
+            mock.expect_shutdown_worker()
+                .times(1)
+                .returning(|_| Err(tonic::Status::unavailable("fake shutdown error")));
+        }
+    } else {
+        mock.expect_shutdown_worker().times(0);
+    }
+
+    let t = canned_histories::single_timer("1");
+    let mut mh =
+        MockPollCfg::from_resp_batches("fakeid", t, [1.into(), ResponseType::AllHistory], mock);
+    mh.enforce_correct_number_of_polls = false;
+    let mut mock = build_mock_pollers(mh);
+    mock.worker_cfg(|w| w.max_cached_workflows = if use_cache { 1 } else { 0 });
+    let worker = mock_worker(mock);
+
+    let res = worker.poll_workflow_activation().await.unwrap();
+    assert_eq!(res.jobs.len(), 1);
+    let run_id = res.run_id;
+
+    tokio::join!(worker.shutdown(), async {
+        // Need to complete task for shutdown to finish
+        worker
+            .complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+                run_id.clone(),
+                workflow_command::Variant::StartTimer(StartTimer {
+                    seq: 1,
+                    start_to_fire_timeout: Some(prost_dur!(from_secs(1))),
+                }),
+            ))
+            .await
+            .unwrap();
+
+        // Shutdown proceeds if the only outstanding activations are evictions
+        assert_matches!(
+            worker.poll_workflow_activation().await.unwrap_err(),
+            PollWfError::ShutDown
+        );
+    });
+}

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

@@ -16,7 +16,7 @@ use crate::{
     },
     Worker,
 };
-use futures::{stream, FutureExt};
+use futures_util::{stream, FutureExt};
 use mockall::TimesRange;
 use rstest::{fixture, rstest};
 use std::{
@@ -33,7 +33,8 @@ use temporal_sdk::{ActivityOptions, CancellableFuture, WfContext};
 use temporal_sdk_core_api::{
     errors::PollWfError,
     worker::{
-        SlotKind, SlotReservationContext, SlotSupplier, SlotSupplierPermit, WorkflowSlotKind,
+        SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext, SlotSupplier,
+        SlotSupplierPermit, WorkflowSlotKind,
     },
     Worker as WorkerTrait,
 };
@@ -2421,9 +2422,7 @@ async fn lang_internal_flag_with_update() {
     let mut t = TestHistoryBuilder::default();
     t.add_by_type(EventType::WorkflowExecutionStarted);
     t.add_full_wf_task();
-    t.set_flags_first_wft(&[1, 2], &[]);
-    t.add_full_wf_task();
-    t.set_flags_last_wft(&[], &[1]);
+    t.set_flags_last_wft(&[1, 2], &[1]);
     let updid = t.add_update_accepted("upd1", "upd");
     t.add_update_completed(updid);
     t.add_workflow_execution_completed();
@@ -3028,8 +3027,8 @@ async fn slot_provider_cant_hand_out_more_permits_than_cache_size() {
         fn try_reserve_slot(&self, _: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
             Some(SlotSupplierPermit::default())
         }
-        fn mark_slot_used(&self, _: <Self::SlotKind as SlotKind>::Info<'_>) {}
-        fn release_slot(&self) {}
+        fn mark_slot_used(&self, _ctx: &dyn SlotMarkUsedContext<SlotKind = Self::SlotKind>) {}
+        fn release_slot(&self, _: &dyn SlotReleaseContext<SlotKind = Self::SlotKind>) {}
         fn available_slots(&self) -> Option<usize> {
             None
         }

package/sdk-core/core/src/debug_client.rs

@@ -0,0 +1,78 @@
+//! Defines an http client that is used for the VSCode debug plugin and any other associated
+//! machinery.
+
+use anyhow::Context;
+use hyper::http::HeaderValue;
+use prost::Message;
+use reqwest::{self, header::HeaderMap};
+use std::time::Duration;
+use temporal_sdk_core_protos::temporal::api::history::v1::History;
+use url::Url;
+
+/// A client for interacting with the VSCode debug plugin
+#[derive(Clone)]
+pub struct DebugClient {
+    /// URL for the local instance of the debugger server
+    debugger_url: Url,
+    client: reqwest::Client,
+}
+
+#[derive(Clone, serde::Serialize)]
+struct WFTStartedMsg {
+    event_id: i64,
+}
+
+impl DebugClient {
+    /// Create a new instance of a DebugClient with the specified url and client name/version
+    /// strings.
+    pub fn new(
+        url: String,
+        client_name: &str,
+        client_version: &str,
+    ) -> Result<DebugClient, anyhow::Error> {
+        let mut client = reqwest::ClientBuilder::new();
+        client = client.default_headers({
+            let mut hm = HeaderMap::new();
+            hm.insert("temporal-client-name", HeaderValue::from_str(client_name)?);
+            hm.insert(
+                "temporal-client-version",
+                HeaderValue::from_str(client_version)?,
+            );
+            hm
+        });
+        let client = client.build()?;
+        Ok(DebugClient {
+            debugger_url: Url::parse(&url).context(
+                "debugger url malformed, is the TEMPORAL_DEBUGGER_PLUGIN_URL env var correct?",
+            )?,
+            client,
+        })
+    }
+
+    /// Get the history from the instance of the debug plugin server
+    pub async fn get_history(&self) -> Result<History, anyhow::Error> {
+        let url = self.debugger_url.join("history")?;
+        let resp = self.client.get(url).send().await?;
+
+        let bytes = resp.bytes().await?;
+        Ok(History::decode(bytes)?)
+    }
+
+    /// Post to current-wft-started to tell the debug plugin which event we've most recently made it
+    /// to
+    pub async fn post_wft_started(
+        &self,
+        event_id: &i64,
+    ) -> Result<reqwest::Response, anyhow::Error> {
+        let url = self.debugger_url.join("current-wft-started")?;
+        Ok(self
+            .client
+            .get(url)
+            .timeout(Duration::from_secs(5))
+            .json(&WFTStartedMsg {
+                event_id: *event_id,
+            })
+            .send()
+            .await?)
+    }
+}

package/sdk-core/core/src/ephemeral_server/mod.rs

@@ -3,7 +3,7 @@
 
 use anyhow::anyhow;
 use flate2::read::GzDecoder;
-use futures::StreamExt;
+use futures_util::StreamExt;
 use serde::Deserialize;
 use std::{
     fs::OpenOptions,
@@ -37,10 +37,13 @@ pub struct TemporalDevServerConfig {
     /// Port to use or obtains a free one if none given.
     #[builder(default)]
     pub port: Option<u16>,
+    /// Port to use for the UI server or obtains a free one if none given.
+    #[builder(default)]
+    pub ui_port: Option<u16>,
     /// Sqlite DB filename if persisting or non-persistent if none.
     #[builder(default)]
     pub db_filename: Option<String>,
-    /// Whether to enable the UI.
+    /// Whether to enable the UI. If ui_port is set, assumes true.
     #[builder(default)]
     pub ui: bool,
     /// Log format and level
@@ -101,7 +104,15 @@ impl TemporalDevServerConfig {
             args.push("--filename".to_owned());
             args.push(db_filename.clone());
         }
-        if !self.ui {
+        if let Some(ui_port) = self.ui_port {
+            args.push("--ui-port".to_owned());
+            args.push(ui_port.to_string());
+        } else if self.ui {
+            // Caps at u16 max which is also max port value
+            let port = port.saturating_add(1000);
+            args.push("--ui-port".to_owned());
+            args.push(port.to_string());
+        } else {
             args.push("--headless".to_owned());
         }
         args.extend(self.extra_args.clone());
@@ -490,7 +501,7 @@ async fn lazy_download_exe(
             info!("Downloading {} to {}", uri, dest.display());
             download_and_extract(client, uri, file_to_extract, &mut temp_file)
                 .await
-                .map_err(|err| {
+                .inspect_err(|_| {
                     // Failed to download, just remove file
                     if let Err(err) = std::fs::remove_file(temp_dest) {
                         warn!(
@@ -499,7 +510,6 @@ async fn lazy_download_exe(
                             err
                         );
                     }
-                    err
                 })
         }
     }?;