Skip to main content

dex_sdk/
client.rs

1// Copyright (c) 2026 Super Durable, Inc.
2//
3// Licensed under the Super Durable Source License 1.0.
4// You may not use this file except in compliance with the License.
5// See the LICENSE file in the repository root.
6//
7// SPDX-License-Identifier: LicenseRef-Super-Durable-1.0
8
9use std::any::TypeId;
10use std::collections::BTreeMap;
11use std::sync::Arc;
12use std::time::{Duration, SystemTime};
13
14use dex_protocol::dex::flow_service_client::FlowServiceClient;
15use dex_protocol::dex::{
16    ActiveStepSearchMode as ProtoSearchMode, AttributeStoreNames, AttributeWrite,
17    FlowAlreadyStartedOptions, FlowConfig as ProtoFlowConfig, FlowResetStepMethod, FlowResetType,
18    FlowRetryPolicy, FlowStartOptions, FlowStatus as ProtoFlowStatus,
19    FlowTimeoutPolicy as ProtoFlowTimeoutPolicy, GetAttributesRequest, GetFlowSummaryRequest,
20    IdReusePolicy as ProtoIdReusePolicy, InvokeRpcRequest, PublishToChannelRequest,
21    ResetFlowRequest, SearchFlowsRequest, SetAttributesRequest, SkipTimerRequest, StartFlowRequest,
22    StepDurability as ProtoStepDurability, StopFlowRequest, StopType as ProtoStopType,
23    TriggerContinueAsNewRequest, UpdateFlowConfigRequest, WaitForAttributeCondition,
24    WaitForAttributeEqual, WaitForAttributeRequest, WaitForFlowRequest,
25    WaitForStepCompletionRequest, WorkerTarget as ProtoWorkerTarget, wait_for_attribute_condition,
26};
27use tokio::runtime::Runtime;
28use tonic::transport::Endpoint;
29use uuid::Uuid;
30
31use crate::sdk_error::{FlowTargetRequirement, ServiceError};
32use crate::stop_flow_options::StopType;
33use crate::time_travel_options::{TimeTravelPoint, TimeTravelStepMethod};
34use crate::value_hydrator::ValueHydrator;
35use crate::value_mapper;
36use crate::worker_dispatcher::map_step_options;
37use crate::{
38    ActiveStepSearchMode, Attribute, AttributeMap, BlobCache, Channel, ChannelMap, ClientOptions,
39    Flow, FlowConfig, FlowErrorType, FlowInfo, FlowResult, FlowStatus, FlowTimeoutPolicy,
40    IdReusePolicy, Registry, RetryPolicy, Rpc, SdkError, SdkResult, SearchFlowEntry,
41    SearchFlowsPage, StartFlowOptions, StepCompletion, StepDurability, StepExecutionId,
42    StopFlowOptions, TimeTravelOptions, TimerId, Value, WorkerTarget,
43};
44
45/// Provides blocking, typed control of registered Dex Flows.
46///
47/// Client methods block the calling thread while an internal Tokio runtime performs gRPC and blob
48/// hydration. The Client owns its transport resources; dropping it releases them. Flow, Step, RPC,
49/// Attribute, and Channel arguments are checked against the [`Registry`].
50///
51/// # Examples
52///
53/// ```no_run
54/// use dex_sdk::{BlobCache, BlobCacheConfig, Client, ClientOptions, Registry};
55/// use std::sync::Arc;
56///
57/// let cache = Arc::new(BlobCache::open(BlobCacheConfig::new(
58///     "/tmp/dex-blobs", 64 * 1024 * 1024, 0,
59/// )?)?);
60/// let client = Client::try_new(Registry::new(), cache, ClientOptions::new())?;
61/// client.health_check()?;
62/// # Ok::<(), Box<dyn std::error::Error>>(())
63/// ```
64pub struct Client {
65    runtime: Runtime,
66    registry: Registry,
67    service: FlowServiceClient<tonic::transport::Channel>,
68    hydrator: ValueHydrator,
69    options: ClientOptions,
70}
71
72impl Client {
73    /// Creates a Client or panics when runtime or endpoint initialization fails.
74    ///
75    /// Prefer [`Self::try_new`] when the application can recover from startup configuration errors.
76    ///
77    /// # Panics
78    ///
79    /// Panics if a Tokio runtime cannot be created or the server address is invalid.
80    pub fn new(registry: Registry, blob_cache: Arc<BlobCache>, options: ClientOptions) -> Self {
81        Self::try_new(registry, blob_cache, options)
82            .unwrap_or_else(|error| panic!("cannot create Dex Client: {error}"))
83    }
84
85    /// Creates a Client from validated definitions, a shared BlobCache, and connection options.
86    ///
87    /// # Errors
88    ///
89    /// Returns [`SdkError::Service`] if runtime or endpoint initialization fails. The transport
90    /// connects lazily, so server availability is checked by the first operation.
91    pub fn try_new(
92        registry: Registry,
93        blob_cache: Arc<BlobCache>,
94        options: ClientOptions,
95    ) -> SdkResult<Self> {
96        let runtime = Runtime::new().map_err(service_error)?;
97        let endpoint = Endpoint::from_shared(endpoint_address(options.server_address_value()))
98            .map_err(service_error)?;
99        let service = {
100            let _runtime_guard = runtime.enter();
101            FlowServiceClient::new(endpoint.connect_lazy())
102        };
103        let hydrator = ValueHydrator::new(service.clone(), blob_cache);
104        Ok(Self {
105            runtime,
106            registry,
107            service,
108            hydrator,
109            options,
110        })
111    }
112
113    /// Starts a registered Flow with server-default options.
114    ///
115    /// Returns the server-assigned run ID. `input` is encoded for the registered starting Step.
116    ///
117    /// # Errors
118    ///
119    /// Returns a definition or mapping error locally, [`SdkError::FlowAlreadyStarted`] for a reuse
120    /// conflict, or another service-backed [`SdkError`] when Dex rejects the request.
121    pub fn start_flow<SomeFlow: Flow>(
122        &self,
123        flow: &SomeFlow,
124        flow_id: &str,
125        input: SomeFlow::StartInput,
126    ) -> SdkResult<String> {
127        self.start_flow_with_options(flow, flow_id, input, StartFlowOptions::new())
128    }
129
130    /// Starts a registered Flow with explicit lifecycle, retry, and initial persistence options.
131    ///
132    /// Returns the server-assigned run ID. A missing request ID is replaced with a generated UUID.
133    ///
134    /// # Errors
135    ///
136    /// Returns [`SdkError`] for invalid IDs, unregistered or mismatched Flow types, encoding
137    /// failures, invalid option durations, reuse conflicts, and service failures.
138    pub fn start_flow_with_options<SomeFlow: Flow>(
139        &self,
140        flow: &SomeFlow,
141        flow_id: &str,
142        input: SomeFlow::StartInput,
143        options: StartFlowOptions,
144    ) -> SdkResult<String> {
145        require_name(flow_id, "Flow ID")?;
146        let registered = self.registry.flow(flow.flow_type())?;
147        if registered.type_id != TypeId::of::<SomeFlow>() {
148            return Err(SdkError::FlowDefinition {
149                message: format!(
150                    "Flow {} does not match the registered Rust type",
151                    flow.flow_type()
152                ),
153            });
154        }
155        let (start_step_type, step_input, step_options) = match &registered.start_step {
156            Some(step) => (
157                step.name.to_string(),
158                Some(value_mapper::encode(&input)?),
159                Some(
160                    map_step_options(
161                        registered,
162                        registered
163                            .handler
164                            .step_options(step.name)
165                            .map_err(sdk_handler_error)?,
166                    )
167                    .map_err(sdk_handler_error)?,
168                ),
169            ),
170            None => (String::new(), None, None),
171        };
172        let flow_timeout_seconds = optional_seconds(options.timeout)?;
173        let flow_timeout_policy =
174            map_flow_timeout_policy(registered, flow_timeout_seconds, options.timeout_policy)?;
175        let request = StartFlowRequest {
176            flow_id: flow_id.to_string(),
177            flow_type: registered.name.to_string(),
178            flow_timeout_seconds,
179            flow_timeout_policy,
180            start_step_type,
181            step_input,
182            step_options,
183            flow_start_options: Some(self.map_start_options(registered, &options)?),
184            request_id: options
185                .request_id
186                .clone()
187                .unwrap_or_else(|| Uuid::new_v4().to_string()),
188        };
189        let mut service = self.service.clone();
190        self.runtime.block_on(async {
191            service
192                .start_flow(request)
193                .await
194                .map(|response| response.into_inner().run_id)
195                .map_err(|status| {
196                    SdkError::from_status(
197                        status,
198                        "start_flow",
199                        Some(flow_id),
200                        FlowTargetRequirement::None,
201                    )
202                })
203        })
204    }
205
206    /// Invokes a registered RPC with typed input and decodes its typed output.
207    ///
208    /// # Errors
209    ///
210    /// Returns [`SdkError::RpcLockConflict`] when locks cannot be acquired, WorkerInvocation for a
211    /// handler failure, FlowNotActive for a terminal Flow, or a mapping/service error.
212    pub fn invoke_rpc<Input: Value, Output: Value>(
213        &self,
214        flow_id: &str,
215        rpc: Rpc<Input, Output>,
216        input: Input,
217    ) -> SdkResult<Output> {
218        self.do_invoke_rpc(flow_id, rpc.name(), &input)
219    }
220
221    /// Invokes a registered no-input RPC and decodes its typed output.
222    ///
223    /// Returns the same errors as [`Self::invoke_rpc`].
224    pub fn invoke_rpc_without_input<Output: Value>(
225        &self,
226        flow_id: &str,
227        rpc: Rpc<(), Output>,
228    ) -> SdkResult<Output> {
229        self.do_invoke_rpc(flow_id, rpc.name(), &())
230    }
231
232    /// Reads one Attribute from the current run.
233    ///
234    /// Returns `Ok(None)` when no value is stored.
235    pub fn get_attribute<T: Value>(
236        &self,
237        flow_id: &str,
238        attribute: &Attribute<T>,
239    ) -> SdkResult<Option<T>> {
240        self.get_attribute_value(flow_id, attribute.name())
241    }
242
243    /// Reads one Attribute-map instance from the current run.
244    ///
245    /// Returns `Ok(None)` when the instance is absent.
246    pub fn get_attribute_map_instance<T: Value>(
247        &self,
248        flow_id: &str,
249        attribute: &AttributeMap<T>,
250        instance: &str,
251    ) -> SdkResult<Option<T>> {
252        self.get_attribute_value(
253            flow_id,
254            &crate::registry::physical_name(attribute.name(), instance),
255        )
256    }
257
258    /// Writes one Attribute on an active Flow.
259    ///
260    /// # Errors
261    ///
262    /// Returns [`SdkError::FlowNotActive`], a value-mapping error, or a service failure.
263    pub fn set_attribute<T: Value>(
264        &self,
265        flow_id: &str,
266        attribute: &Attribute<T>,
267        value: T,
268    ) -> SdkResult<()> {
269        self.set_attribute_value(
270            flow_id,
271            attribute.name(),
272            &value,
273            attribute.index().map(|index| index.proto_config(false)),
274            attribute.sync_config(),
275        )
276    }
277
278    /// Writes one Attribute-map instance on an active Flow.
279    ///
280    /// # Errors
281    ///
282    /// Returns [`SdkError::FlowNotActive`], a value-mapping error, or a service failure.
283    pub fn set_attribute_map_instance<T: Value>(
284        &self,
285        flow_id: &str,
286        attribute: &AttributeMap<T>,
287        instance: &str,
288        value: T,
289    ) -> SdkResult<()> {
290        self.set_attribute_value(
291            flow_id,
292            &crate::registry::physical_name(attribute.name(), instance),
293            &value,
294            attribute.index().map(|index| index.proto_config(true)),
295            attribute.sync_config(),
296        )
297    }
298
299    /// Publishes one typed message to a Channel on an active Flow.
300    pub fn publish<T: Value>(
301        &self,
302        flow_id: &str,
303        channel: &Channel<T>,
304        value: T,
305    ) -> SdkResult<()> {
306        self.publish_values(flow_id, channel.name(), [value])
307    }
308
309    /// Publishes an ordered batch of typed messages to one Channel atomically.
310    pub fn publish_many<T: Value>(
311        &self,
312        flow_id: &str,
313        channel: &Channel<T>,
314        values: impl IntoIterator<Item = T>,
315    ) -> SdkResult<()> {
316        self.publish_values(flow_id, channel.name(), values)
317    }
318
319    /// Publishes an ordered batch to one Channel-map instance atomically.
320    pub fn publish_map<T: Value>(
321        &self,
322        flow_id: &str,
323        channel: &ChannelMap<T>,
324        instance: &str,
325        values: impl IntoIterator<Item = T>,
326    ) -> SdkResult<()> {
327        self.publish_values(
328            flow_id,
329            &crate::registry::physical_name(channel.name(), instance),
330            values,
331        )
332    }
333
334    /// Blocks until a Flow closes and returns its terminal result.
335    ///
336    /// # Errors
337    ///
338    /// Returns FlowNotFound for an unknown ID or a mapping/service error. Every terminal status is
339    /// represented by [`FlowResult`]. This overload has no client-side timeout.
340    pub fn wait_for_flow(&self, flow_id: &str) -> SdkResult<FlowResult> {
341        self.wait_for_flow_result(flow_id, None)
342    }
343
344    /// Blocks up to `timeout` and returns all output-bearing completions for a successful Flow.
345    ///
346    /// # Errors
347    ///
348    /// Returns [`SdkError::LongPollTimeout`] when the timeout elapses while the Flow remains active;
349    /// other errors match [`Self::wait_for_flow`].
350    pub fn wait_for_flow_with_timeout(
351        &self,
352        flow_id: &str,
353        timeout: Duration,
354    ) -> SdkResult<FlowResult> {
355        self.wait_for_flow_result(flow_id, Some(timeout))
356    }
357
358    /// Returns current identity, status, type, and start time for the latest run.
359    ///
360    /// # Errors
361    ///
362    /// Returns [`SdkError::FlowNotFound`] for an unknown ID or a service/protocol error.
363    pub fn describe_flow(&self, flow_id: &str) -> SdkResult<FlowInfo> {
364        let mut service = self.service.clone();
365        let response = self.runtime.block_on(async {
366            service
367                .get_flow_summary(GetFlowSummaryRequest {
368                    flow_id: flow_id.to_string(),
369                    run_id: String::new(),
370                })
371                .await
372                .map(|response| response.into_inner())
373                .map_err(|status| {
374                    SdkError::from_status(
375                        status,
376                        "describe_flow",
377                        Some(flow_id),
378                        FlowTargetRequirement::Existing,
379                    )
380                })
381        })?;
382        let execution = response
383            .flow_execution_id
384            .ok_or_else(|| service_error("GetFlowSummary omitted FlowExecutionID"))?;
385        Ok(FlowInfo {
386            flow_id: execution.flow_id,
387            run_id: execution.run_id,
388            flow_type: response.flow_type,
389            status: map_flow_status(response.flow_status)?,
390            started_at: response
391                .start_time
392                .map_or(SystemTime::UNIX_EPOCH, timestamp),
393        })
394    }
395
396    /// Executes a server search query and returns its first page.
397    ///
398    /// `page_size` must be non-negative; zero uses the server default. Indexed values are hydrated
399    /// and decoded as [`serde_json::Value`].
400    pub fn search_flows(&self, query: &str, page_size: i32) -> SdkResult<SearchFlowsPage> {
401        self.search_flows_page(query, page_size, "")
402    }
403
404    /// Continues a search with the opaque token returned by a previous page.
405    ///
406    /// Keep `query` and `page_size` consistent across pages. An empty returned token marks the last
407    /// page.
408    pub fn search_flows_page(
409        &self,
410        query: &str,
411        page_size: i32,
412        next_page_token: &str,
413    ) -> SdkResult<SearchFlowsPage> {
414        if page_size < 0 {
415            return Err(invalid("search page size must not be negative"));
416        }
417        let mut service = self.service.clone();
418        self.runtime.block_on(async {
419            let response = service
420                .search_flows(SearchFlowsRequest {
421                    query: query.to_string(),
422                    page_size,
423                    next_page_token: next_page_token.to_string(),
424                })
425                .await
426                .map(|response| response.into_inner())
427                .map_err(|status| {
428                    SdkError::from_status(status, "search_flows", None, FlowTargetRequirement::None)
429                })?;
430            let values = response
431                .flow_runs
432                .iter()
433                .flat_map(|entry| entry.indexed_attributes.iter())
434                .map(|attribute| {
435                    attribute
436                        .value
437                        .clone()
438                        .ok_or_else(|| invalid("Indexed Attribute has no Value"))
439                })
440                .collect::<SdkResult<Vec<_>>>()?;
441            let mut values = self.hydrator.hydrate_all(values).await?.into_iter();
442            let mut flows = Vec::with_capacity(response.flow_runs.len());
443            for entry in response.flow_runs {
444                let mut indexed_attributes = BTreeMap::new();
445                for attribute in entry.indexed_attributes {
446                    let value = values
447                        .next()
448                        .ok_or_else(|| invalid("Indexed Attribute hydration count mismatch"))?;
449                    indexed_attributes.insert(attribute.key, value_mapper::decode_untyped(&value)?);
450                }
451                flows.push(SearchFlowEntry {
452                    flow_id: entry.flow_id,
453                    run_id: entry.run_id,
454                    flow_type: entry.flow_type,
455                    status: map_flow_status(entry.flow_status)?,
456                    started_at: entry.start_time.map(timestamp),
457                    closed_at: entry.close_time.map(timestamp),
458                    indexed_attributes,
459                });
460            }
461            Ok(SearchFlowsPage {
462                flows,
463                next_page_token: response.next_page_token,
464            })
465        })
466    }
467
468    /// Cancels, terminates, or fails an active Flow according to `options`.
469    ///
470    /// # Errors
471    ///
472    /// Returns [`SdkError::FlowNotActive`] when no active run exists or another service error.
473    pub fn stop_flow(&self, flow_id: &str, options: StopFlowOptions) -> SdkResult<()> {
474        let stop_type = match options.stop_type {
475            StopType::Cancel => ProtoStopType::Cancel,
476            StopType::Terminate => ProtoStopType::Terminate,
477            StopType::Fail => ProtoStopType::Fail,
478        };
479        self.call_empty(
480            "stop_flow",
481            Some(flow_id),
482            FlowTargetRequirement::Active,
483            |mut service| async move {
484                service
485                    .stop_flow(StopFlowRequest {
486                        flow_id: flow_id.to_string(),
487                        run_id: String::new(),
488                        reason: options.reason.unwrap_or_default(),
489                        stop_type: stop_type as i32,
490                    })
491                    .await
492            },
493        )
494    }
495
496    /// Creates a new run by replaying an existing Flow from the selected historical point.
497    ///
498    /// Returns the new server-assigned run ID.
499    ///
500    /// # Errors
501    ///
502    /// Returns InvalidArgument for an unrepresentable time, FlowNotFound for an unknown ID, or
503    /// another service error.
504    pub fn time_travel(&self, flow_id: &str, options: TimeTravelOptions) -> SdkResult<String> {
505        let mut request = ResetFlowRequest {
506            flow_id: flow_id.to_string(),
507            run_id: String::new(),
508            reset_type: FlowResetType::Unspecified as i32,
509            reason: options.reason.unwrap_or_default(),
510            history_event_time: String::new(),
511            step_type: String::new(),
512            step_execution_id: String::new(),
513            skip_writes_reapply: options.skip_writes_reapply,
514            step_method: FlowResetStepMethod::Unspecified as i32,
515        };
516        match options.point {
517            TimeTravelPoint::Beginning => request.reset_type = FlowResetType::Beginning as i32,
518            TimeTravelPoint::HistoryEventTime(time) => {
519                request.reset_type = FlowResetType::HistoryEventTime as i32;
520                request.history_event_time = rfc3339(time)?;
521            }
522            TimeTravelPoint::StepType(step_type) => {
523                request.reset_type = FlowResetType::StepType as i32;
524                request.step_type = step_type.to_string();
525            }
526            TimeTravelPoint::StepExecution(execution, method) => {
527                request.reset_type = FlowResetType::StepExecutionId as i32;
528                request.step_execution_id =
529                    format!("{}-{}", execution.step_type, execution.execution_number);
530                request.step_method = match method {
531                    TimeTravelStepMethod::WaitFor => FlowResetStepMethod::WaitFor as i32,
532                    TimeTravelStepMethod::Execute => FlowResetStepMethod::Execute as i32,
533                };
534            }
535        }
536        let mut service = self.service.clone();
537        self.runtime.block_on(async {
538            service
539                .reset_flow(request)
540                .await
541                .map(|response| response.into_inner().run_id)
542                .map_err(|status| {
543                    SdkError::from_status(
544                        status,
545                        "time_travel",
546                        Some(flow_id),
547                        FlowTargetRequirement::Existing,
548                    )
549                })
550        })
551    }
552
553    /// Marks one timer in an active Step execution as satisfied immediately.
554    ///
555    /// # Errors
556    ///
557    /// Returns InvalidArgument for an oversized index, FlowNotActive for a terminal Flow, or a
558    /// service error when the Step execution or timer cannot be targeted.
559    pub fn skip_timer(
560        &self,
561        flow_id: &str,
562        step_execution: StepExecutionId,
563        timer: TimerId,
564    ) -> SdkResult<()> {
565        let (timer_condition_id, timer_condition_index) = match timer {
566            TimerId::ConditionId(id) => (id, None),
567            TimerId::ConditionIndex(index) => (
568                String::new(),
569                Some(i32::try_from(index).map_err(|_| invalid("timer index exceeds int32"))?),
570            ),
571        };
572        self.call_empty(
573            "skip_timer",
574            Some(flow_id),
575            FlowTargetRequirement::Active,
576            |mut service| async move {
577                service
578                    .skip_timer(SkipTimerRequest {
579                        flow_id: flow_id.to_string(),
580                        run_id: String::new(),
581                        step_execution_id: format!(
582                            "{}-{}",
583                            step_execution.step_type, step_execution.execution_number
584                        ),
585                        timer_condition_id,
586                        timer_condition_index,
587                    })
588                    .await
589            },
590        )
591    }
592
593    /// Blocks until one Step execution completes or `timeout` elapses.
594    ///
595    /// # Errors
596    ///
597    /// Returns [`SdkError::LongPollTimeout`] on timeout, FlowNotActive when appropriate, or another
598    /// service error. Successful completion returns `()` and does not decode Step output.
599    pub fn wait_for_step_completion(
600        &self,
601        flow_id: &str,
602        step_execution: StepExecutionId,
603        timeout: Duration,
604    ) -> SdkResult<()> {
605        let wait_time_seconds = seconds32(timeout)?;
606        self.call_empty(
607            "wait_for_step_completion",
608            Some(flow_id),
609            FlowTargetRequirement::Active,
610            |mut service| async move {
611                service
612                    .wait_for_step_completion(WaitForStepCompletionRequest {
613                        flow_id: flow_id.to_string(),
614                        step_type: step_execution.step_type.to_string(),
615                        step_execution_number: step_execution.execution_number.to_string(),
616                        wait_time_seconds,
617                        request_id: Uuid::new_v4().to_string(),
618                    })
619                    .await
620            },
621        )
622    }
623
624    /// Blocks until a singleton Attribute in the current run equals `expected`.
625    ///
626    /// String, bool, integer, and double values are supported. Object, bytes,
627    /// and null values return [`SdkError::InvalidArgument`] before transport.
628    /// A server-side expiry returns [`SdkError::LongPollTimeout`].
629    pub fn wait_for_attribute_equal<T: Value>(
630        &self,
631        flow_id: &str,
632        attribute: &Attribute<T>,
633        expected: T,
634        timeout: Duration,
635    ) -> SdkResult<()> {
636        self.wait_for_attribute_value(flow_id, attribute.name(), &expected, timeout)
637    }
638
639    /// Blocks until one AttributeMap instance equals `expected`.
640    ///
641    /// This targets the current run and otherwise has the same primitive-value,
642    /// timeout, request-ID, and error behavior as `wait_for_attribute_equal`.
643    pub fn wait_for_attribute_map_instance_equal<T: Value>(
644        &self,
645        flow_id: &str,
646        attribute: &AttributeMap<T>,
647        instance: &str,
648        expected: T,
649        timeout: Duration,
650    ) -> SdkResult<()> {
651        self.wait_for_attribute_value(
652            flow_id,
653            &crate::registry::physical_name(attribute.name(), instance),
654            &expected,
655            timeout,
656        )
657    }
658
659    fn wait_for_attribute_value<T: Value>(
660        &self,
661        flow_id: &str,
662        key: &str,
663        expected: &T,
664        timeout: Duration,
665    ) -> SdkResult<()> {
666        let value = value_mapper::encode(expected)?;
667        if !matches!(
668            value.kind,
669            Some(dex_protocol::dex::value::Kind::StringValue(_))
670                | Some(dex_protocol::dex::value::Kind::BoolValue(_))
671                | Some(dex_protocol::dex::value::Kind::IntValue(_))
672                | Some(dex_protocol::dex::value::Kind::DoubleValue(_))
673        ) {
674            return Err(invalid(
675                "wait_for_attribute_equal supports only string, boolean, or number values",
676            ));
677        }
678        let wait_time_seconds = seconds32(timeout)?;
679        self.call_empty(
680            "wait_for_attribute_equal",
681            Some(flow_id),
682            FlowTargetRequirement::Active,
683            |mut service| async move {
684                service
685                    .wait_for_attribute(WaitForAttributeRequest {
686                        flow_id: flow_id.to_string(),
687                        run_id: String::new(),
688                        condition: Some(WaitForAttributeCondition {
689                            kind: Some(wait_for_attribute_condition::Kind::Equal(
690                                WaitForAttributeEqual {
691                                    key: key.to_string(),
692                                    value: Some(value),
693                                },
694                            )),
695                        }),
696                        wait_time_seconds,
697                        request_id: Uuid::new_v4().to_string(),
698                    })
699                    .await
700            },
701        )
702    }
703
704    /// Replaces mutable runtime configuration fields on an active Flow.
705    ///
706    /// Unset fields retain server-defined behavior.
707    pub fn update_flow_config(&self, flow_id: &str, config: FlowConfig) -> SdkResult<()> {
708        let flow_config = self.map_flow_config(Some(&config))?;
709        self.call_empty(
710            "update_flow_config",
711            Some(flow_id),
712            FlowTargetRequirement::Active,
713            |mut service| async move {
714                service
715                    .update_flow_config(UpdateFlowConfigRequest {
716                        flow_id: flow_id.to_string(),
717                        run_id: String::new(),
718                        flow_config: Some(flow_config),
719                    })
720                    .await
721            },
722        )
723    }
724
725    /// Requests that an active Flow roll its history into a successor run.
726    pub fn trigger_continue_as_new(&self, flow_id: &str) -> SdkResult<()> {
727        self.call_empty(
728            "trigger_continue_as_new",
729            Some(flow_id),
730            FlowTargetRequirement::Active,
731            |mut service| async move {
732                service
733                    .trigger_continue_as_new(TriggerContinueAsNewRequest {
734                        flow_id: flow_id.to_string(),
735                        run_id: String::new(),
736                    })
737                    .await
738            },
739        )
740    }
741
742    /// Verifies that Dex FlowService responds successfully.
743    ///
744    /// Returns `Ok(())` for a healthy service and a service-backed [`SdkError`] otherwise.
745    pub fn health_check(&self) -> SdkResult<()> {
746        self.call_empty(
747            "health_check",
748            None,
749            FlowTargetRequirement::None,
750            |mut service| async move { service.health_check(()).await },
751        )
752    }
753
754    fn do_invoke_rpc<Input: Value, Output: Value>(
755        &self,
756        flow_id: &str,
757        rpc_name: &str,
758        input: &Input,
759    ) -> SdkResult<Output> {
760        let rpc = self.registry.rpc(rpc_name)?;
761        let request = InvokeRpcRequest {
762            flow_id: flow_id.to_string(),
763            run_id: String::new(),
764            rpc_name: rpc_name.to_string(),
765            input: Some(value_mapper::encode(input)?),
766            timeout_seconds: optional_seconds(rpc.timeout)?,
767            lock_attribute_keys: rpc.locks.iter().map(|lock| lock.physical_name()).collect(),
768            request_id: Uuid::new_v4().to_string(),
769        };
770        let mut service = self.service.clone();
771        let output = self.runtime.block_on(async {
772            service
773                .invoke_rpc(request)
774                .await
775                .map(|response| response.into_inner().output)
776                .map_err(|status| {
777                    SdkError::from_status(
778                        status,
779                        "invoke_rpc",
780                        Some(flow_id),
781                        FlowTargetRequirement::Active,
782                    )
783                })
784        })?;
785        let output = output.ok_or_else(|| invalid("InvokeRPC omitted output"))?;
786        let output = self.runtime.block_on(self.hydrator.hydrate(output))?;
787        value_mapper::decode(&output)
788    }
789
790    fn get_attribute_value<T: Value>(&self, flow_id: &str, key: &str) -> SdkResult<Option<T>> {
791        let mut service = self.service.clone();
792        let response = self.runtime.block_on(async {
793            service
794                .get_attributes(GetAttributesRequest {
795                    flow_id: flow_id.to_string(),
796                    run_id: String::new(),
797                    keys: vec![key.to_string()],
798                    all_keys: false,
799                })
800                .await
801                .map(|response| response.into_inner())
802                .map_err(|status| {
803                    SdkError::from_status(
804                        status,
805                        "get_attribute",
806                        Some(flow_id),
807                        FlowTargetRequirement::Existing,
808                    )
809                })
810        })?;
811        let Some(entry) = response.attributes.into_iter().next() else {
812            return Ok(None);
813        };
814        let value = entry
815            .value
816            .ok_or_else(|| invalid("GetAttributes returned an empty Value"))?;
817        let value = self.runtime.block_on(self.hydrator.hydrate(value))?;
818        value_mapper::decode(&value).map(Some)
819    }
820
821    fn set_attribute_value<T: Value>(
822        &self,
823        flow_id: &str,
824        key: &str,
825        value: &T,
826        index_config: Option<dex_protocol::dex::IndexConfig>,
827        sync_config: Option<dex_protocol::dex::AttributeSyncConfig>,
828    ) -> SdkResult<()> {
829        let write = AttributeWrite {
830            key: key.to_string(),
831            value: Some(value_mapper::encode(value)?),
832            index_config,
833            sync_config,
834        };
835        self.call_empty(
836            "set_attribute",
837            Some(flow_id),
838            FlowTargetRequirement::Active,
839            |mut service| async move {
840                service
841                    .set_attributes(SetAttributesRequest {
842                        flow_id: flow_id.to_string(),
843                        run_id: String::new(),
844                        attributes: vec![write],
845                        request_id: Uuid::new_v4().to_string(),
846                    })
847                    .await
848            },
849        )
850    }
851
852    fn publish_values<T: Value>(
853        &self,
854        flow_id: &str,
855        channel_name: &str,
856        values: impl IntoIterator<Item = T>,
857    ) -> SdkResult<()> {
858        let messages = values
859            .into_iter()
860            .map(|value| {
861                Ok(dex_protocol::dex::ChannelMessage {
862                    channel_name: channel_name.to_string(),
863                    value: Some(value_mapper::encode(&value)?),
864                })
865            })
866            .collect::<SdkResult<Vec<_>>>()?;
867        self.call_empty(
868            "publish",
869            Some(flow_id),
870            FlowTargetRequirement::Active,
871            |mut service| async move {
872                service
873                    .publish_to_channel(PublishToChannelRequest {
874                        flow_id: flow_id.to_string(),
875                        run_id: String::new(),
876                        messages,
877                    })
878                    .await
879            },
880        )
881    }
882
883    fn wait_for_flow_result(
884        &self,
885        flow_id: &str,
886        timeout: Option<Duration>,
887    ) -> SdkResult<FlowResult> {
888        let mut service = self.service.clone();
889        let response = self.runtime.block_on(async {
890            service
891                .wait_for_flow(WaitForFlowRequest {
892                    flow_id: flow_id.to_string(),
893                    run_id: String::new(),
894                    needs_results: true,
895                    wait_time_seconds: optional_seconds(timeout)?,
896                })
897                .await
898                .map(|response| response.into_inner())
899                .map_err(|status| {
900                    SdkError::from_status(
901                        status,
902                        "wait_for_flow",
903                        Some(flow_id),
904                        FlowTargetRequirement::Existing,
905                    )
906                })
907        });
908        let response = response?;
909        let mut outputs = Vec::with_capacity(response.results.len());
910        for result in &response.results {
911            outputs.push(result.completed_step_output.clone().ok_or_else(|| {
912                SdkError::ValueMapping {
913                    message: "Step completion output is required".to_string(),
914                }
915            })?);
916        }
917        let outputs = self.runtime.block_on(self.hydrator.hydrate_all(outputs))?;
918        let completions = response
919            .results
920            .into_iter()
921            .zip(outputs)
922            .map(|(result, output)| {
923                StepCompletion::new(
924                    result.completed_step_type,
925                    result.completed_step_execution_id,
926                    output,
927                )
928            })
929            .collect();
930        Ok(FlowResult::new(
931            map_flow_status(response.flow_status)?,
932            map_flow_error_type(response.error_type),
933            (!response.error_message.is_empty()).then_some(response.error_message),
934            completions,
935        ))
936    }
937
938    fn map_start_options(
939        &self,
940        flow: &crate::registry::RegisteredFlow,
941        options: &StartFlowOptions,
942    ) -> SdkResult<FlowStartOptions> {
943        let attributes = options
944            .attributes
945            .iter()
946            .map(|attribute| {
947                let logical_name = attribute.key.split('/').next().unwrap_or(&attribute.key);
948                if !flow.persistence.contains_key(logical_name) {
949                    return Err(invalid(format!(
950                        "initial Attribute is not registered: {}",
951                        attribute.key
952                    )));
953                }
954                Ok(AttributeWrite {
955                    key: attribute.key.clone(),
956                    value: Some(attribute.value.encode()?),
957                    index_config: attribute.index_config.clone(),
958                    sync_config: attribute.sync_config,
959                })
960            })
961            .collect::<SdkResult<Vec<_>>>()?;
962        let config = options.config_override.as_ref();
963        let flow_config_override =
964            if config.is_some() || self.options.worker_target_value().is_some() {
965                Some(self.map_flow_config(config)?)
966            } else {
967                None
968            };
969        #[allow(clippy::needless_update)]
970        let flow_start_options = FlowStartOptions {
971            id_reuse_policy: match options.id_reuse_policy {
972                IdReusePolicy::Default => ProtoIdReusePolicy::Unspecified,
973                IdReusePolicy::AllowIfPreviousFailed => {
974                    ProtoIdReusePolicy::AllowIfPreviousExistsAbnormally
975                }
976                IdReusePolicy::AllowIfNotRunning => ProtoIdReusePolicy::AllowIfNoRunning,
977                IdReusePolicy::Disallow => ProtoIdReusePolicy::DisallowReuse,
978                IdReusePolicy::TerminateIfRunning => ProtoIdReusePolicy::AllowTerminateIfRunning,
979            } as i32,
980            flow_start_delay_seconds: optional_seconds(options.start_delay)?,
981            retry_policy: options
982                .retry_policy
983                .clone()
984                .map(map_flow_retry)
985                .transpose()?,
986            attributes,
987            flow_config_override,
988            flow_already_started_options: Some(FlowAlreadyStartedOptions {
989                ignore_already_started_error: options.ignore_already_started,
990            }),
991            ..Default::default()
992        };
993        Ok(flow_start_options)
994    }
995
996    fn map_flow_config(&self, config: Option<&FlowConfig>) -> SdkResult<ProtoFlowConfig> {
997        map_flow_config(config, self.options.worker_target_value())
998    }
999
1000    fn call_empty<Response, Future, Call>(
1001        &self,
1002        operation: &'static str,
1003        flow_id: Option<&str>,
1004        requirement: FlowTargetRequirement,
1005        call: Call,
1006    ) -> SdkResult<()>
1007    where
1008        Call: FnOnce(FlowServiceClient<tonic::transport::Channel>) -> Future,
1009        Future: std::future::Future<Output = Result<tonic::Response<Response>, tonic::Status>>,
1010    {
1011        let service = self.service.clone();
1012        self.runtime
1013            .block_on(call(service))
1014            .map(|_| ())
1015            .map_err(|status| SdkError::from_status(status, operation, flow_id, requirement))
1016    }
1017}
1018
1019pub(crate) fn map_flow_config(
1020    config: Option<&FlowConfig>,
1021    default_target: Option<&WorkerTarget>,
1022) -> SdkResult<ProtoFlowConfig> {
1023    let target = config
1024        .and_then(|config| config.worker_target.as_ref())
1025        .or(default_target);
1026    Ok(ProtoFlowConfig {
1027        active_step_search_mode: config
1028            .and_then(|config| config.active_step_search_mode)
1029            .map(|mode| match mode {
1030                ActiveStepSearchMode::All => ProtoSearchMode::EnabledForAll as i32,
1031                ActiveStepSearchMode::WithWaitFor => {
1032                    ProtoSearchMode::EnabledForStepsWithWaitFor as i32
1033                }
1034                ActiveStepSearchMode::Disabled => ProtoSearchMode::Disabled as i32,
1035            }),
1036        continue_as_new_threshold: config
1037            .and_then(|config| config.continue_as_new_threshold)
1038            .map(i32::try_from)
1039            .transpose()
1040            .map_err(|_| invalid("continue-as-new threshold exceeds int32"))?,
1041        continue_as_new_page_size_in_bytes: config
1042            .and_then(|config| config.continue_as_new_page_size_bytes)
1043            .map(i32::try_from)
1044            .transpose()
1045            .map_err(|_| invalid("continue-as-new page size exceeds int32"))?,
1046        step_durability: config
1047            .and_then(|config| config.step_durability)
1048            .map(|durability| {
1049                (match durability {
1050                    StepDurability::Default => ProtoStepDurability::Unspecified,
1051                    StepDurability::Sync => ProtoStepDurability::Sync,
1052                    StepDurability::Async => ProtoStepDurability::Async,
1053                }) as i32
1054            }),
1055        worker_target: target.map(map_worker_target),
1056        attribute_store_names: config.and_then(|config| {
1057            config
1058                .attribute_store_names
1059                .as_ref()
1060                .map(|names| AttributeStoreNames {
1061                    names: names.clone(),
1062                })
1063        }),
1064    })
1065}
1066
1067pub(crate) fn map_flow_retry(retry: RetryPolicy) -> SdkResult<FlowRetryPolicy> {
1068    Ok(FlowRetryPolicy {
1069        initial_interval_seconds: optional_seconds(retry.initial_interval)?,
1070        backoff_coefficient: retry.backoff_coefficient.unwrap_or_default() as f32,
1071        maximum_interval_seconds: optional_seconds(retry.maximum_interval)?,
1072        maximum_attempts: i32::try_from(retry.maximum_attempts.unwrap_or_default())
1073            .map_err(|_| invalid("maximum attempts exceed int32"))?,
1074    })
1075}
1076
1077fn map_worker_target(target: &WorkerTarget) -> ProtoWorkerTarget {
1078    ProtoWorkerTarget {
1079        address: target.address().to_string(),
1080        is_headless_address: target.is_headless(),
1081    }
1082}
1083
1084pub(crate) fn map_flow_status(status: i32) -> SdkResult<FlowStatus> {
1085    match ProtoFlowStatus::try_from(status).ok() {
1086        Some(ProtoFlowStatus::Running) => Ok(FlowStatus::Running),
1087        Some(ProtoFlowStatus::Completed) => Ok(FlowStatus::Completed),
1088        Some(ProtoFlowStatus::Failed) => Ok(FlowStatus::Failed),
1089        Some(ProtoFlowStatus::ServerSideTimeoutInternalOnly) => {
1090            Ok(FlowStatus::ServerSideTimeoutInternalOnly)
1091        }
1092        Some(ProtoFlowStatus::Terminated) => Ok(FlowStatus::Terminated),
1093        Some(ProtoFlowStatus::Canceled) => Ok(FlowStatus::Canceled),
1094        Some(ProtoFlowStatus::ContinuedAsNew) => Ok(FlowStatus::ContinuedAsNew),
1095        _ => Err(invalid(format!("unknown Flow status {status}"))),
1096    }
1097}
1098
1099pub(crate) fn map_flow_error_type(error_type: i32) -> Option<FlowErrorType> {
1100    use dex_protocol::dex::FlowErrorType as ProtoFlowErrorType;
1101
1102    match ProtoFlowErrorType::try_from(error_type).ok() {
1103        Some(ProtoFlowErrorType::StepDecisionFailingFlow) => {
1104            Some(FlowErrorType::StepDecisionFailed)
1105        }
1106        Some(ProtoFlowErrorType::ClientApiFailingFlow) => Some(FlowErrorType::ClientApiFailed),
1107        Some(ProtoFlowErrorType::WorkerApiFail) => Some(FlowErrorType::WorkerApiFailed),
1108        Some(ProtoFlowErrorType::InvalidUserFlowCode) => Some(FlowErrorType::InvalidUserFlowCode),
1109        Some(ProtoFlowErrorType::Internal) => Some(FlowErrorType::Internal),
1110        Some(ProtoFlowErrorType::FlowTimeout) => Some(FlowErrorType::FlowTimeout),
1111        _ => None,
1112    }
1113}
1114
1115pub(crate) fn map_flow_timeout_policy(
1116    flow: &crate::registry::RegisteredFlow,
1117    timeout_seconds: i32,
1118    policy: FlowTimeoutPolicy,
1119) -> SdkResult<i32> {
1120    if timeout_seconds == 0 {
1121        if policy != FlowTimeoutPolicy::Default {
1122            return Err(invalid("Flow timeout policy requires a positive timeout"));
1123        }
1124        return Ok(ProtoFlowTimeoutPolicy::Unspecified as i32);
1125    }
1126    let resolved = match policy {
1127        FlowTimeoutPolicy::Default if flow.handler.has_timeout_handler() => {
1128            FlowTimeoutPolicy::Handler
1129        }
1130        FlowTimeoutPolicy::Default => FlowTimeoutPolicy::Fail,
1131        configured => configured,
1132    };
1133    if resolved == FlowTimeoutPolicy::Handler && !flow.handler.has_timeout_handler() {
1134        return Err(invalid(format!(
1135            "Flow {} has no timeout handler",
1136            flow.name
1137        )));
1138    }
1139    Ok(match resolved {
1140        FlowTimeoutPolicy::Default => ProtoFlowTimeoutPolicy::Unspecified,
1141        FlowTimeoutPolicy::Fail => ProtoFlowTimeoutPolicy::Fail,
1142        FlowTimeoutPolicy::Cancel => ProtoFlowTimeoutPolicy::Cancel,
1143        FlowTimeoutPolicy::Handler => ProtoFlowTimeoutPolicy::Handler,
1144    } as i32)
1145}
1146
1147fn timestamp(timestamp: prost_types::Timestamp) -> SystemTime {
1148    if timestamp.seconds >= 0 {
1149        SystemTime::UNIX_EPOCH
1150            + Duration::from_secs(timestamp.seconds as u64)
1151            + Duration::from_nanos(timestamp.nanos.max(0) as u64)
1152    } else {
1153        SystemTime::UNIX_EPOCH - Duration::from_secs(timestamp.seconds.unsigned_abs())
1154    }
1155}
1156
1157fn rfc3339(time: SystemTime) -> SdkResult<String> {
1158    let duration = time
1159        .duration_since(SystemTime::UNIX_EPOCH)
1160        .map_err(service_error)?;
1161    let days = (duration.as_secs() / 86_400) as i64;
1162    let seconds = duration.as_secs() % 86_400;
1163    let (year, month, day) = civil_from_days(days);
1164    let hour = seconds / 3_600;
1165    let minute = seconds % 3_600 / 60;
1166    let second = seconds % 60;
1167    Ok(format!(
1168        "{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z"
1169    ))
1170}
1171
1172fn civil_from_days(days: i64) -> (i64, i64, i64) {
1173    let days = days + 719_468;
1174    let era = if days >= 0 { days } else { days - 146_096 } / 146_097;
1175    let day_of_era = days - era * 146_097;
1176    let year_of_era =
1177        (day_of_era - day_of_era / 1_460 + day_of_era / 36_524 - day_of_era / 146_096) / 365;
1178    let mut year = year_of_era + era * 400;
1179    let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100);
1180    let month_prime = (5 * day_of_year + 2) / 153;
1181    let day = day_of_year - (153 * month_prime + 2) / 5 + 1;
1182    let month = month_prime + if month_prime < 10 { 3 } else { -9 };
1183    year += i64::from(month <= 2);
1184    (year, month, day)
1185}
1186
1187fn optional_seconds(duration: Option<Duration>) -> SdkResult<i32> {
1188    duration
1189        .map(seconds32)
1190        .transpose()
1191        .map(Option::unwrap_or_default)
1192}
1193
1194fn seconds32(duration: Duration) -> SdkResult<i32> {
1195    if duration.subsec_nanos() != 0 {
1196        return Err(invalid("Duration must use whole seconds"));
1197    }
1198    i32::try_from(duration.as_secs()).map_err(|_| invalid("Duration exceeds int32"))
1199}
1200
1201fn endpoint_address(address: &str) -> String {
1202    if address.contains("://") {
1203        address.to_string()
1204    } else {
1205        format!("http://{address}")
1206    }
1207}
1208
1209fn require_name(value: &str, kind: &str) -> SdkResult<()> {
1210    if value.is_empty() {
1211        Err(invalid(format!("{kind} is required")))
1212    } else {
1213        Ok(())
1214    }
1215}
1216
1217fn sdk_handler_error(error: impl std::fmt::Display) -> SdkError {
1218    invalid(error.to_string())
1219}
1220
1221fn invalid(message: impl Into<String>) -> SdkError {
1222    SdkError::InvalidArgument {
1223        message: message.into(),
1224    }
1225}
1226
1227fn service_error(error: impl std::fmt::Display) -> SdkError {
1228    SdkError::Service {
1229        service: ServiceError::local("client", error.to_string()),
1230    }
1231}
1232
1233#[cfg(test)]
1234mod tests {
1235    use super::map_flow_config;
1236    use crate::FlowConfig;
1237
1238    #[test]
1239    fn attribute_store_names_preserve_protocol_presence() {
1240        let absent = map_flow_config(Some(&FlowConfig::new()), None).expect("map absent config");
1241        let named = map_flow_config(
1242            Some(
1243                &FlowConfig::new()
1244                    .attribute_store_names(vec!["profiles".to_owned(), "audit".to_owned()]),
1245            ),
1246            None,
1247        )
1248        .expect("map named config");
1249        let disabled =
1250            map_flow_config(Some(&FlowConfig::new().attribute_store_names(vec![])), None)
1251                .expect("map disabled config");
1252
1253        assert_eq!(absent.attribute_store_names, None);
1254        assert_eq!(
1255            named.attribute_store_names.expect("named stores").names,
1256            vec!["profiles", "audit"]
1257        );
1258        assert_eq!(
1259            disabled
1260                .attribute_store_names
1261                .expect("disabled stores")
1262                .names,
1263            Vec::<String>::new()
1264        );
1265    }
1266}