dex_sdk/

client.rs

// Copyright (c) 2026 Super Durable, Inc.
//
// Licensed under the Super Durable Source License 1.0.
// You may not use this file except in compliance with the License.
// See the LICENSE file in the repository root.
//
// SPDX-License-Identifier: LicenseRef-Super-Durable-1.0

use std::any::TypeId;
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime};

use dex_protocol::dex::flow_service_client::FlowServiceClient;
use dex_protocol::dex::{
    ActiveStepSearchMode as ProtoSearchMode, AttributeStoreNames, AttributeWrite,
    FlowAlreadyStartedOptions, FlowConfig as ProtoFlowConfig, FlowResetStepMethod, FlowResetType,
    FlowRetryPolicy, FlowStartOptions, FlowStatus as ProtoFlowStatus,
    FlowTimeoutPolicy as ProtoFlowTimeoutPolicy, GetAttributesRequest, GetFlowSummaryRequest,
    IdReusePolicy as ProtoIdReusePolicy, InvokeRpcRequest, PublishToChannelRequest,
    ReadStreamRequest, ResetFlowRequest, SearchFlowsRequest, SetAttributesRequest,
    SkipTimerRequest, StartFlowRequest, StepDurability as ProtoStepDurability, StopFlowRequest,
    StopType as ProtoStopType, TriggerContinueAsNewRequest, UpdateFlowConfigRequest,
    WaitForAttributeCondition, WaitForAttributeEqual, WaitForAttributeRequest, WaitForFlowRequest,
    WaitForStepCompletionRequest, WorkerTarget as ProtoWorkerTarget, WriteStreamRequest,
    wait_for_attribute_condition,
};
use tokio::runtime::Runtime;
use tonic::transport::Endpoint;
use uuid::Uuid;

use crate::sdk_error::{FlowTargetRequirement, ServiceError};
use crate::stop_flow_options::StopType;
use crate::time_travel_options::{TimeTravelPoint, TimeTravelStepMethod};
use crate::value_hydrator::ValueHydrator;
use crate::value_mapper;
use crate::worker_dispatcher::map_step_options;
use crate::{
    ActiveStepSearchMode, Attribute, AttributeMap, BlobCache, Channel, ChannelMap, ClientOptions,
    Flow, FlowConfig, FlowErrorType, FlowInfo, FlowResult, FlowStatus, FlowTimeoutPolicy,
    IdReusePolicy, Registry, RetryPolicy, Rpc, SdkError, SdkResult, SearchFlowEntry,
    SearchFlowsPage, StartFlowOptions, StepCompletion, StepDurability, StepExecutionId,
    StopFlowOptions, Stream, StreamMessage, TimeTravelOptions, TimerId, Value, WorkerTarget,
};

/// Provides blocking, typed control of registered Dex Flows.
///
/// Client methods block the calling thread while an internal Tokio runtime performs gRPC and blob
/// hydration. The Client owns its transport resources; dropping it releases them. Flow, Step, RPC,
/// Attribute, and Channel arguments are checked against the [`Registry`].
///
/// # Examples
///
/// ```no_run
/// use dex_sdk::{BlobCache, BlobCacheConfig, Client, ClientOptions, Registry};
/// use std::sync::Arc;
///
/// let cache = Arc::new(BlobCache::open(BlobCacheConfig::new(
///     "/tmp/dex-blobs", 64 * 1024 * 1024, 0,
/// )?)?);
/// let client = Client::try_new(Registry::new(), cache, ClientOptions::new())?;
/// client.health_check()?;
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
pub struct Client {
    runtime: Runtime,
    registry: Registry,
    service: FlowServiceClient<tonic::transport::Channel>,
    hydrator: ValueHydrator,
    options: ClientOptions,
}

impl Client {
    /// Creates a Client or panics when runtime or endpoint initialization fails.
    ///
    /// Prefer [`Self::try_new`] when the application can recover from startup configuration errors.
    ///
    /// # Panics
    ///
    /// Panics if a Tokio runtime cannot be created or the server address is invalid.
    pub fn new(registry: Registry, blob_cache: Arc<BlobCache>, options: ClientOptions) -> Self {
        Self::try_new(registry, blob_cache, options)
            .unwrap_or_else(|error| panic!("cannot create Dex Client: {error}"))
    }

    /// Creates a Client from validated definitions, a shared BlobCache, and connection options.
    ///
    /// # Errors
    ///
    /// Returns [`SdkError::Service`] if runtime or endpoint initialization fails. The transport
    /// connects lazily, so server availability is checked by the first operation.
    pub fn try_new(
        registry: Registry,
        blob_cache: Arc<BlobCache>,
        options: ClientOptions,
    ) -> SdkResult<Self> {
        let runtime = Runtime::new().map_err(service_error)?;
        let endpoint = Endpoint::from_shared(endpoint_address(options.server_address_value()))
            .map_err(service_error)?;
        let service = {
            let _runtime_guard = runtime.enter();
            FlowServiceClient::new(endpoint.connect_lazy())
        };
        let hydrator = ValueHydrator::new(service.clone(), blob_cache);
        Ok(Self {
            runtime,
            registry,
            service,
            hydrator,
            options,
        })
    }

    /// Starts a registered Flow with server-default options.
    ///
    /// Returns the server-assigned run ID. `input` is encoded for the registered starting Step.
    ///
    /// # Errors
    ///
    /// Returns a definition or mapping error locally, [`SdkError::FlowAlreadyStarted`] for a reuse
    /// conflict, or another service-backed [`SdkError`] when Dex rejects the request.
    pub fn start_flow<SomeFlow: Flow>(
        &self,
        flow: &SomeFlow,
        flow_id: &str,
        input: SomeFlow::StartInput,
    ) -> SdkResult<String> {
        self.start_flow_with_options(flow, flow_id, input, StartFlowOptions::new())
    }

    /// Starts a registered Flow with explicit lifecycle, retry, and initial persistence options.
    ///
    /// Returns the server-assigned run ID. A missing request ID is replaced with a generated UUID.
    ///
    /// # Errors
    ///
    /// Returns [`SdkError`] for invalid IDs, unregistered or mismatched Flow types, encoding
    /// failures, invalid option durations, reuse conflicts, and service failures.
    pub fn start_flow_with_options<SomeFlow: Flow>(
        &self,
        flow: &SomeFlow,
        flow_id: &str,
        input: SomeFlow::StartInput,
        options: StartFlowOptions,
    ) -> SdkResult<String> {
        require_name(flow_id, "Flow ID")?;
        let registered = self.registry.flow(flow.flow_type())?;
        if registered.type_id != TypeId::of::<SomeFlow>() {
            return Err(SdkError::FlowDefinition {
                message: format!(
                    "Flow {} does not match the registered Rust type",
                    flow.flow_type()
                ),
            });
        }
        let (start_step_type, step_input, step_options) = match &registered.start_step {
            Some(step) => (
                step.name.to_string(),
                Some(value_mapper::encode(&input)?),
                Some(
                    map_step_options(
                        registered,
                        registered
                            .handler
                            .step_options(step.name)
                            .map_err(sdk_handler_error)?,
                    )
                    .map_err(sdk_handler_error)?,
                ),
            ),
            None => (String::new(), None, None),
        };
        let flow_timeout_seconds = optional_seconds(options.timeout)?;
        let flow_timeout_policy =
            map_flow_timeout_policy(registered, flow_timeout_seconds, options.timeout_policy)?;
        let request = StartFlowRequest {
            flow_id: flow_id.to_string(),
            flow_type: registered.name.to_string(),
            flow_timeout_seconds,
            flow_timeout_policy,
            start_step_type,
            step_input,
            step_options,
            flow_start_options: Some(self.map_start_options(registered, &options)?),
            request_id: options
                .request_id
                .clone()
                .unwrap_or_else(|| Uuid::new_v4().to_string()),
        };
        let mut service = self.service.clone();
        self.runtime.block_on(async {
            service
                .start_flow(request)
                .await
                .map(|response| response.into_inner().run_id)
                .map_err(|status| {
                    SdkError::from_status(
                        status,
                        "start_flow",
                        Some(flow_id),
                        FlowTargetRequirement::None,
                    )
                })
        })
    }

    /// Invokes a registered RPC with typed input and decodes its typed output.
    ///
    /// # Errors
    ///
    /// Returns [`SdkError::RpcLockConflict`] when locks cannot be acquired, WorkerInvocation for a
    /// handler failure, FlowNotActive for a terminal Flow, or a mapping/service error.
    pub fn invoke_rpc<Input: Value, Output: Value>(
        &self,
        flow_id: &str,
        rpc: Rpc<Input, Output>,
        input: Input,
    ) -> SdkResult<Output> {
        self.do_invoke_rpc(flow_id, rpc.name(), &input)
    }

    /// Invokes a registered no-input RPC and decodes its typed output.
    ///
    /// Returns the same errors as [`Self::invoke_rpc`].
    pub fn invoke_rpc_without_input<Output: Value>(
        &self,
        flow_id: &str,
        rpc: Rpc<(), Output>,
    ) -> SdkResult<Output> {
        self.do_invoke_rpc(flow_id, rpc.name(), &())
    }

    /// Reads one Attribute from the current run.
    ///
    /// Returns `Ok(None)` when no value is stored.
    pub fn get_attribute<T: Value>(
        &self,
        flow_id: &str,
        attribute: &Attribute<T>,
    ) -> SdkResult<Option<T>> {
        self.get_attribute_value(flow_id, attribute.name())
    }

    /// Reads one Attribute-map instance from the current run.
    ///
    /// Returns `Ok(None)` when the instance is absent.
    pub fn get_attribute_map_instance<T: Value>(
        &self,
        flow_id: &str,
        attribute: &AttributeMap<T>,
        instance: &str,
    ) -> SdkResult<Option<T>> {
        self.get_attribute_value(
            flow_id,
            &crate::registry::physical_name(attribute.name(), instance),
        )
    }

    /// Writes one Attribute on an active Flow.
    ///
    /// # Errors
    ///
    /// Returns [`SdkError::FlowNotActive`], a value-mapping error, or a service failure.
    pub fn set_attribute<T: Value>(
        &self,
        flow_id: &str,
        attribute: &Attribute<T>,
        value: T,
    ) -> SdkResult<()> {
        self.set_attribute_value(
            flow_id,
            attribute.name(),
            &value,
            attribute.index().map(|index| index.proto_config(false)),
            attribute.sync_config(),
        )
    }

    /// Writes one Attribute-map instance on an active Flow.
    ///
    /// # Errors
    ///
    /// Returns [`SdkError::FlowNotActive`], a value-mapping error, or a service failure.
    pub fn set_attribute_map_instance<T: Value>(
        &self,
        flow_id: &str,
        attribute: &AttributeMap<T>,
        instance: &str,
        value: T,
    ) -> SdkResult<()> {
        self.set_attribute_value(
            flow_id,
            &crate::registry::physical_name(attribute.name(), instance),
            &value,
            attribute.index().map(|index| index.proto_config(true)),
            attribute.sync_config(),
        )
    }

    /// Publishes one typed message to a Channel on an active Flow.
    pub fn publish<T: Value>(
        &self,
        flow_id: &str,
        channel: &Channel<T>,
        value: T,
    ) -> SdkResult<()> {
        self.publish_values(flow_id, channel.name(), [value])
    }

    /// Publishes an ordered batch of typed messages to one Channel atomically.
    pub fn publish_many<T: Value>(
        &self,
        flow_id: &str,
        channel: &Channel<T>,
        values: impl IntoIterator<Item = T>,
    ) -> SdkResult<()> {
        self.publish_values(flow_id, channel.name(), values)
    }

    /// Publishes an ordered batch to one Channel-map instance atomically.
    pub fn publish_map<T: Value>(
        &self,
        flow_id: &str,
        channel: &ChannelMap<T>,
        instance: &str,
        values: impl IntoIterator<Item = T>,
    ) -> SdkResult<()> {
        self.publish_values(
            flow_id,
            &crate::registry::physical_name(channel.name(), instance),
            values,
        )
    }

    /// Appends one typed best-effort Stream message with client idempotency.
    ///
    /// The Flow instance need not exist or be active. Retained duplicate keys are successful
    /// first-write-wins no-ops.
    ///
    /// # Errors
    ///
    /// Returns a definition error for an unregistered Stream, InvalidArgument for an empty ID or
    /// key containing `#`, a mapping error, or a FlowService failure.
    pub fn write_stream<T: Value>(
        &self,
        flow_id: &str,
        stream: &Stream<T>,
        idempotency_key: &str,
        value: T,
    ) -> SdkResult<()> {
        require_name(flow_id, "Flow ID")?;
        require_name(idempotency_key, "Stream idempotency key")?;
        if idempotency_key.contains('#') {
            return Err(invalid("Stream client idempotency key must not contain #"));
        }
        let flow_type = self.registry.flow_for_stream(stream)?.name.to_string();
        let request = WriteStreamRequest {
            flow_id: flow_id.to_string(),
            flow_type,
            stream_name: stream.name().to_string(),
            max_estimated_bytes: stream.max_estimated_bytes(),
            value: Some(value_mapper::encode(&value)?),
            idempotency_key: idempotency_key.to_string(),
        };
        self.call_empty(
            "write_stream",
            Some(flow_id),
            FlowTargetRequirement::None,
            |mut service| async move { service.write_stream(request).await },
        )
    }

    /// Blocks for the next retained Stream message using the server default wait.
    ///
    /// An empty token starts at the retained head. Pass the returned token unchanged to resume.
    pub fn read_stream<T: Value>(
        &self,
        flow_id: &str,
        stream: &Stream<T>,
        resume_token: &str,
    ) -> SdkResult<StreamMessage<T>> {
        self.read_stream_result(flow_id, stream, resume_token, None)
    }

    /// Blocks for the next retained Stream message with an explicit server wait.
    ///
    /// # Errors
    ///
    /// Returns [`SdkError::LongPollTimeout`] when no message arrives, plus the same definition,
    /// duration, mapping, and service errors as [`Self::read_stream`].
    pub fn read_stream_with_timeout<T: Value>(
        &self,
        flow_id: &str,
        stream: &Stream<T>,
        resume_token: &str,
        timeout: Duration,
    ) -> SdkResult<StreamMessage<T>> {
        self.read_stream_result(flow_id, stream, resume_token, Some(timeout))
    }

    /// Blocks until a Flow closes and returns its terminal result.
    ///
    /// # Errors
    ///
    /// Returns FlowNotFound for an unknown ID or a mapping/service error. Every terminal status is
    /// represented by [`FlowResult`]. This overload has no client-side timeout.
    pub fn wait_for_flow(&self, flow_id: &str) -> SdkResult<FlowResult> {
        self.wait_for_flow_result(flow_id, None)
    }

    /// Blocks up to `timeout` and returns all output-bearing completions for a successful Flow.
    ///
    /// # Errors
    ///
    /// Returns [`SdkError::LongPollTimeout`] when the timeout elapses while the Flow remains active;
    /// other errors match [`Self::wait_for_flow`].
    pub fn wait_for_flow_with_timeout(
        &self,
        flow_id: &str,
        timeout: Duration,
    ) -> SdkResult<FlowResult> {
        self.wait_for_flow_result(flow_id, Some(timeout))
    }

    /// Returns current identity, status, type, and start time for the latest run.
    ///
    /// # Errors
    ///
    /// Returns [`SdkError::FlowNotFound`] for an unknown ID or a service/protocol error.
    pub fn describe_flow(&self, flow_id: &str) -> SdkResult<FlowInfo> {
        let mut service = self.service.clone();
        let response = self.runtime.block_on(async {
            service
                .get_flow_summary(GetFlowSummaryRequest {
                    flow_id: flow_id.to_string(),
                    run_id: String::new(),
                })
                .await
                .map(|response| response.into_inner())
                .map_err(|status| {
                    SdkError::from_status(
                        status,
                        "describe_flow",
                        Some(flow_id),
                        FlowTargetRequirement::Existing,
                    )
                })
        })?;
        let execution = response
            .flow_execution_id
            .ok_or_else(|| service_error("GetFlowSummary omitted FlowExecutionID"))?;
        Ok(FlowInfo {
            flow_id: execution.flow_id,
            run_id: execution.run_id,
            flow_type: response.flow_type,
            status: map_flow_status(response.flow_status)?,
            started_at: response
                .start_time
                .map_or(SystemTime::UNIX_EPOCH, timestamp),
        })
    }

    /// Executes a server search query and returns its first page.
    ///
    /// `page_size` must be non-negative; zero uses the server default. Indexed values are hydrated
    /// and decoded as [`serde_json::Value`].
    pub fn search_flows(&self, query: &str, page_size: i32) -> SdkResult<SearchFlowsPage> {
        self.search_flows_page(query, page_size, "")
    }

    /// Continues a search with the opaque token returned by a previous page.
    ///
    /// Keep `query` and `page_size` consistent across pages. An empty returned token marks the last
    /// page.
    pub fn search_flows_page(
        &self,
        query: &str,
        page_size: i32,
        next_page_token: &str,
    ) -> SdkResult<SearchFlowsPage> {
        if page_size < 0 {
            return Err(invalid("search page size must not be negative"));
        }
        let mut service = self.service.clone();
        self.runtime.block_on(async {
            let response = service
                .search_flows(SearchFlowsRequest {
                    query: query.to_string(),
                    page_size,
                    next_page_token: next_page_token.to_string(),
                })
                .await
                .map(|response| response.into_inner())
                .map_err(|status| {
                    SdkError::from_status(status, "search_flows", None, FlowTargetRequirement::None)
                })?;
            let values = response
                .flow_runs
                .iter()
                .flat_map(|entry| entry.indexed_attributes.iter())
                .map(|attribute| {
                    attribute
                        .value
                        .clone()
                        .ok_or_else(|| invalid("Indexed Attribute has no Value"))
                })
                .collect::<SdkResult<Vec<_>>>()?;
            let mut values = self.hydrator.hydrate_all(values).await?.into_iter();
            let mut flows = Vec::with_capacity(response.flow_runs.len());
            for entry in response.flow_runs {
                let mut indexed_attributes = BTreeMap::new();
                for attribute in entry.indexed_attributes {
                    let value = values
                        .next()
                        .ok_or_else(|| invalid("Indexed Attribute hydration count mismatch"))?;
                    indexed_attributes.insert(attribute.key, value_mapper::decode_untyped(&value)?);
                }
                flows.push(SearchFlowEntry {
                    flow_id: entry.flow_id,
                    run_id: entry.run_id,
                    flow_type: entry.flow_type,
                    status: map_flow_status(entry.flow_status)?,
                    started_at: entry.start_time.map(timestamp),
                    closed_at: entry.close_time.map(timestamp),
                    indexed_attributes,
                });
            }
            Ok(SearchFlowsPage {
                flows,
                next_page_token: response.next_page_token,
            })
        })
    }

    /// Cancels, terminates, or fails an active Flow according to `options`.
    ///
    /// # Errors
    ///
    /// Returns [`SdkError::FlowNotActive`] when no active run exists or another service error.
    pub fn stop_flow(&self, flow_id: &str, options: StopFlowOptions) -> SdkResult<()> {
        let stop_type = match options.stop_type {
            StopType::Cancel => ProtoStopType::Cancel,
            StopType::Terminate => ProtoStopType::Terminate,
            StopType::Fail => ProtoStopType::Fail,
        };
        self.call_empty(
            "stop_flow",
            Some(flow_id),
            FlowTargetRequirement::Active,
            |mut service| async move {
                service
                    .stop_flow(StopFlowRequest {
                        flow_id: flow_id.to_string(),
                        run_id: String::new(),
                        reason: options.reason.unwrap_or_default(),
                        stop_type: stop_type as i32,
                    })
                    .await
            },
        )
    }

    /// Creates a new run by replaying an existing Flow from the selected historical point.
    ///
    /// Returns the new server-assigned run ID.
    ///
    /// # Errors
    ///
    /// Returns InvalidArgument for an unrepresentable time, FlowNotFound for an unknown ID, or
    /// another service error.
    pub fn time_travel(&self, flow_id: &str, options: TimeTravelOptions) -> SdkResult<String> {
        let mut request = ResetFlowRequest {
            flow_id: flow_id.to_string(),
            run_id: String::new(),
            reset_type: FlowResetType::Unspecified as i32,
            reason: options.reason.unwrap_or_default(),
            history_event_time: String::new(),
            step_type: String::new(),
            step_execution_id: String::new(),
            skip_writes_reapply: options.skip_writes_reapply,
            step_method: FlowResetStepMethod::Unspecified as i32,
        };
        match options.point {
            TimeTravelPoint::Beginning => request.reset_type = FlowResetType::Beginning as i32,
            TimeTravelPoint::HistoryEventTime(time) => {
                request.reset_type = FlowResetType::HistoryEventTime as i32;
                request.history_event_time = rfc3339(time)?;
            }
            TimeTravelPoint::StepType(step_type) => {
                request.reset_type = FlowResetType::StepType as i32;
                request.step_type = step_type.to_string();
            }
            TimeTravelPoint::StepExecution(execution, method) => {
                request.reset_type = FlowResetType::StepExecutionId as i32;
                request.step_execution_id =
                    format!("{}-{}", execution.step_type, execution.execution_number);
                request.step_method = match method {
                    TimeTravelStepMethod::WaitFor => FlowResetStepMethod::WaitFor as i32,
                    TimeTravelStepMethod::Execute => FlowResetStepMethod::Execute as i32,
                };
            }
        }
        let mut service = self.service.clone();
        self.runtime.block_on(async {
            service
                .reset_flow(request)
                .await
                .map(|response| response.into_inner().run_id)
                .map_err(|status| {
                    SdkError::from_status(
                        status,
                        "time_travel",
                        Some(flow_id),
                        FlowTargetRequirement::Existing,
                    )
                })
        })
    }

    /// Marks one timer in an active Step execution as satisfied immediately.
    ///
    /// # Errors
    ///
    /// Returns InvalidArgument for an oversized index, FlowNotActive for a terminal Flow, or a
    /// service error when the Step execution or timer cannot be targeted.
    pub fn skip_timer(
        &self,
        flow_id: &str,
        step_execution: StepExecutionId,
        timer: TimerId,
    ) -> SdkResult<()> {
        let (timer_condition_id, timer_condition_index) = match timer {
            TimerId::ConditionId(id) => (id, None),
            TimerId::ConditionIndex(index) => (
                String::new(),
                Some(i32::try_from(index).map_err(|_| invalid("timer index exceeds int32"))?),
            ),
        };
        self.call_empty(
            "skip_timer",
            Some(flow_id),
            FlowTargetRequirement::Active,
            |mut service| async move {
                service
                    .skip_timer(SkipTimerRequest {
                        flow_id: flow_id.to_string(),
                        run_id: String::new(),
                        step_execution_id: format!(
                            "{}-{}",
                            step_execution.step_type, step_execution.execution_number
                        ),
                        timer_condition_id,
                        timer_condition_index,
                    })
                    .await
            },
        )
    }

    /// Blocks until one Step execution completes or `timeout` elapses.
    ///
    /// # Errors
    ///
    /// Returns [`SdkError::LongPollTimeout`] on timeout, FlowNotActive when appropriate, or another
    /// service error. Successful completion returns `()` and does not decode Step output.
    pub fn wait_for_step_completion(
        &self,
        flow_id: &str,
        step_execution: StepExecutionId,
        timeout: Duration,
    ) -> SdkResult<()> {
        let wait_time_seconds = seconds32(timeout)?;
        self.call_empty(
            "wait_for_step_completion",
            Some(flow_id),
            FlowTargetRequirement::Active,
            |mut service| async move {
                service
                    .wait_for_step_completion(WaitForStepCompletionRequest {
                        flow_id: flow_id.to_string(),
                        step_type: step_execution.step_type.to_string(),
                        step_execution_number: step_execution.execution_number.to_string(),
                        wait_time_seconds,
                        request_id: Uuid::new_v4().to_string(),
                    })
                    .await
            },
        )
    }

    /// Blocks until a singleton Attribute in the current run equals `expected`.
    ///
    /// String, bool, integer, and double values are supported. Object, bytes,
    /// and null values return [`SdkError::InvalidArgument`] before transport.
    /// A server-side expiry returns [`SdkError::LongPollTimeout`].
    pub fn wait_for_attribute_equal<T: Value>(
        &self,
        flow_id: &str,
        attribute: &Attribute<T>,
        expected: T,
        timeout: Duration,
    ) -> SdkResult<()> {
        self.wait_for_attribute_value(flow_id, attribute.name(), &expected, timeout)
    }

    /// Blocks until one AttributeMap instance equals `expected`.
    ///
    /// This targets the current run and otherwise has the same primitive-value,
    /// timeout, request-ID, and error behavior as `wait_for_attribute_equal`.
    pub fn wait_for_attribute_map_instance_equal<T: Value>(
        &self,
        flow_id: &str,
        attribute: &AttributeMap<T>,
        instance: &str,
        expected: T,
        timeout: Duration,
    ) -> SdkResult<()> {
        self.wait_for_attribute_value(
            flow_id,
            &crate::registry::physical_name(attribute.name(), instance),
            &expected,
            timeout,
        )
    }

    fn wait_for_attribute_value<T: Value>(
        &self,
        flow_id: &str,
        key: &str,
        expected: &T,
        timeout: Duration,
    ) -> SdkResult<()> {
        let value = value_mapper::encode(expected)?;
        if !matches!(
            value.kind,
            Some(dex_protocol::dex::value::Kind::StringValue(_))
                | Some(dex_protocol::dex::value::Kind::BoolValue(_))
                | Some(dex_protocol::dex::value::Kind::IntValue(_))
                | Some(dex_protocol::dex::value::Kind::DoubleValue(_))
        ) {
            return Err(invalid(
                "wait_for_attribute_equal supports only string, boolean, or number values",
            ));
        }
        let wait_time_seconds = seconds32(timeout)?;
        self.call_empty(
            "wait_for_attribute_equal",
            Some(flow_id),
            FlowTargetRequirement::Active,
            |mut service| async move {
                service
                    .wait_for_attribute(WaitForAttributeRequest {
                        flow_id: flow_id.to_string(),
                        run_id: String::new(),
                        condition: Some(WaitForAttributeCondition {
                            kind: Some(wait_for_attribute_condition::Kind::Equal(
                                WaitForAttributeEqual {
                                    key: key.to_string(),
                                    value: Some(value),
                                },
                            )),
                        }),
                        wait_time_seconds,
                        request_id: Uuid::new_v4().to_string(),
                    })
                    .await
            },
        )
    }

    /// Replaces mutable runtime configuration fields on an active Flow.
    ///
    /// Unset fields retain server-defined behavior.
    pub fn update_flow_config(&self, flow_id: &str, config: FlowConfig) -> SdkResult<()> {
        let flow_config = self.map_flow_config(Some(&config))?;
        self.call_empty(
            "update_flow_config",
            Some(flow_id),
            FlowTargetRequirement::Active,
            |mut service| async move {
                service
                    .update_flow_config(UpdateFlowConfigRequest {
                        flow_id: flow_id.to_string(),
                        run_id: String::new(),
                        flow_config: Some(flow_config),
                    })
                    .await
            },
        )
    }

    /// Requests that an active Flow roll its history into a successor run.
    pub fn trigger_continue_as_new(&self, flow_id: &str) -> SdkResult<()> {
        self.call_empty(
            "trigger_continue_as_new",
            Some(flow_id),
            FlowTargetRequirement::Active,
            |mut service| async move {
                service
                    .trigger_continue_as_new(TriggerContinueAsNewRequest {
                        flow_id: flow_id.to_string(),
                        run_id: String::new(),
                    })
                    .await
            },
        )
    }

    /// Verifies that Dex FlowService responds successfully.
    ///
    /// Returns `Ok(())` for a healthy service and a service-backed [`SdkError`] otherwise.
    pub fn health_check(&self) -> SdkResult<()> {
        self.call_empty(
            "health_check",
            None,
            FlowTargetRequirement::None,
            |mut service| async move { service.health_check(()).await },
        )
    }

    fn do_invoke_rpc<Input: Value, Output: Value>(
        &self,
        flow_id: &str,
        rpc_name: &str,
        input: &Input,
    ) -> SdkResult<Output> {
        let rpc = self.registry.rpc(rpc_name)?;
        let request = InvokeRpcRequest {
            flow_id: flow_id.to_string(),
            run_id: String::new(),
            rpc_name: rpc_name.to_string(),
            input: Some(value_mapper::encode(input)?),
            timeout_seconds: optional_seconds(rpc.timeout)?,
            lock_attribute_keys: rpc.locks.iter().map(|lock| lock.physical_name()).collect(),
            request_id: Uuid::new_v4().to_string(),
        };
        let mut service = self.service.clone();
        let output = self.runtime.block_on(async {
            service
                .invoke_rpc(request)
                .await
                .map(|response| response.into_inner().output)
                .map_err(|status| {
                    SdkError::from_status(
                        status,
                        "invoke_rpc",
                        Some(flow_id),
                        FlowTargetRequirement::Active,
                    )
                })
        })?;
        let output = output.ok_or_else(|| invalid("InvokeRPC omitted output"))?;
        let output = self.runtime.block_on(self.hydrator.hydrate(output))?;
        value_mapper::decode(&output)
    }

    fn get_attribute_value<T: Value>(&self, flow_id: &str, key: &str) -> SdkResult<Option<T>> {
        let mut service = self.service.clone();
        let response = self.runtime.block_on(async {
            service
                .get_attributes(GetAttributesRequest {
                    flow_id: flow_id.to_string(),
                    run_id: String::new(),
                    keys: vec![key.to_string()],
                    all_keys: false,
                })
                .await
                .map(|response| response.into_inner())
                .map_err(|status| {
                    SdkError::from_status(
                        status,
                        "get_attribute",
                        Some(flow_id),
                        FlowTargetRequirement::Existing,
                    )
                })
        })?;
        let Some(entry) = response.attributes.into_iter().next() else {
            return Ok(None);
        };
        let value = entry
            .value
            .ok_or_else(|| invalid("GetAttributes returned an empty Value"))?;
        let value = self.runtime.block_on(self.hydrator.hydrate(value))?;
        value_mapper::decode(&value).map(Some)
    }

    fn set_attribute_value<T: Value>(
        &self,
        flow_id: &str,
        key: &str,
        value: &T,
        index_config: Option<dex_protocol::dex::IndexConfig>,
        sync_config: Option<dex_protocol::dex::AttributeSyncConfig>,
    ) -> SdkResult<()> {
        let write = AttributeWrite {
            key: key.to_string(),
            value: Some(value_mapper::encode(value)?),
            index_config,
            sync_config,
        };
        self.call_empty(
            "set_attribute",
            Some(flow_id),
            FlowTargetRequirement::Active,
            |mut service| async move {
                service
                    .set_attributes(SetAttributesRequest {
                        flow_id: flow_id.to_string(),
                        run_id: String::new(),
                        attributes: vec![write],
                        request_id: Uuid::new_v4().to_string(),
                    })
                    .await
            },
        )
    }

    fn publish_values<T: Value>(
        &self,
        flow_id: &str,
        channel_name: &str,
        values: impl IntoIterator<Item = T>,
    ) -> SdkResult<()> {
        let messages = values
            .into_iter()
            .map(|value| {
                Ok(dex_protocol::dex::ChannelMessage {
                    channel_name: channel_name.to_string(),
                    value: Some(value_mapper::encode(&value)?),
                })
            })
            .collect::<SdkResult<Vec<_>>>()?;
        self.call_empty(
            "publish",
            Some(flow_id),
            FlowTargetRequirement::Active,
            |mut service| async move {
                service
                    .publish_to_channel(PublishToChannelRequest {
                        flow_id: flow_id.to_string(),
                        run_id: String::new(),
                        messages,
                    })
                    .await
            },
        )
    }

    fn read_stream_result<T: Value>(
        &self,
        flow_id: &str,
        stream: &Stream<T>,
        resume_token: &str,
        timeout: Option<Duration>,
    ) -> SdkResult<StreamMessage<T>> {
        require_name(flow_id, "Flow ID")?;
        let flow_type = self.registry.flow_for_stream(stream)?.name.to_string();
        let wait_time_seconds = optional_seconds(timeout)?;
        let request = ReadStreamRequest {
            flow_id: flow_id.to_string(),
            flow_type,
            stream_name: stream.name().to_string(),
            resume_token: resume_token.to_string(),
            wait_time_seconds,
        };
        let mut service = self.service.clone();
        let response = self.runtime.block_on(async {
            service
                .read_stream(request)
                .await
                .map(|response| response.into_inner())
                .map_err(|status| {
                    SdkError::from_status(
                        status,
                        "read_stream",
                        Some(flow_id),
                        FlowTargetRequirement::None,
                    )
                })
        })?;
        let message = response
            .message
            .ok_or_else(|| invalid("ReadStream omitted its message"))?;
        if message.resume_token.is_empty() {
            return Err(invalid("ReadStream returned an empty resume token"));
        }
        let value = message
            .value
            .ok_or_else(|| invalid("ReadStream omitted its Value"))?;
        let created_time = message
            .created_time
            .map(timestamp)
            .ok_or_else(|| invalid("ReadStream omitted its creation time"))?;
        Ok(StreamMessage {
            value: value_mapper::decode(&value)?,
            resume_token: message.resume_token,
            created_time,
            idempotency_key: message.idempotency_key,
        })
    }

    fn wait_for_flow_result(
        &self,
        flow_id: &str,
        timeout: Option<Duration>,
    ) -> SdkResult<FlowResult> {
        let mut service = self.service.clone();
        let response = self.runtime.block_on(async {
            service
                .wait_for_flow(WaitForFlowRequest {
                    flow_id: flow_id.to_string(),
                    run_id: String::new(),
                    needs_results: true,
                    wait_time_seconds: optional_seconds(timeout)?,
                })
                .await
                .map(|response| response.into_inner())
                .map_err(|status| {
                    SdkError::from_status(
                        status,
                        "wait_for_flow",
                        Some(flow_id),
                        FlowTargetRequirement::Existing,
                    )
                })
        });
        let response = response?;
        let mut outputs = Vec::with_capacity(response.results.len());
        for result in &response.results {
            outputs.push(result.completed_step_output.clone().ok_or_else(|| {
                SdkError::ValueMapping {
                    message: "Step completion output is required".to_string(),
                }
            })?);
        }
        let outputs = self.runtime.block_on(self.hydrator.hydrate_all(outputs))?;
        let completions = response
            .results
            .into_iter()
            .zip(outputs)
            .map(|(result, output)| {
                StepCompletion::new(
                    result.completed_step_type,
                    result.completed_step_execution_id,
                    output,
                )
            })
            .collect();
        Ok(FlowResult::new(
            map_flow_status(response.flow_status)?,
            map_flow_error_type(response.error_type),
            (!response.error_message.is_empty()).then_some(response.error_message),
            completions,
        ))
    }

    fn map_start_options(
        &self,
        flow: &crate::registry::RegisteredFlow,
        options: &StartFlowOptions,
    ) -> SdkResult<FlowStartOptions> {
        let attributes = options
            .attributes
            .iter()
            .map(|attribute| {
                let logical_name = attribute.key.split('/').next().unwrap_or(&attribute.key);
                if !flow.persistence.contains_key(logical_name) {
                    return Err(invalid(format!(
                        "initial Attribute is not registered: {}",
                        attribute.key
                    )));
                }
                Ok(AttributeWrite {
                    key: attribute.key.clone(),
                    value: Some(attribute.value.encode()?),
                    index_config: attribute.index_config.clone(),
                    sync_config: attribute.sync_config,
                })
            })
            .collect::<SdkResult<Vec<_>>>()?;
        let config = options.config_override.as_ref();
        let flow_config_override =
            if config.is_some() || self.options.worker_target_value().is_some() {
                Some(self.map_flow_config(config)?)
            } else {
                None
            };
        #[allow(clippy::needless_update)]
        let flow_start_options = FlowStartOptions {
            id_reuse_policy: match options.id_reuse_policy {
                IdReusePolicy::Default => ProtoIdReusePolicy::Unspecified,
                IdReusePolicy::AllowIfPreviousFailed => {
                    ProtoIdReusePolicy::AllowIfPreviousExistsAbnormally
                }
                IdReusePolicy::AllowIfNotRunning => ProtoIdReusePolicy::AllowIfNoRunning,
                IdReusePolicy::Disallow => ProtoIdReusePolicy::DisallowReuse,
                IdReusePolicy::TerminateIfRunning => ProtoIdReusePolicy::AllowTerminateIfRunning,
            } as i32,
            flow_start_delay_seconds: optional_seconds(options.start_delay)?,
            retry_policy: options
                .retry_policy
                .clone()
                .map(map_flow_retry)
                .transpose()?,
            attributes,
            flow_config_override,
            flow_already_started_options: Some(FlowAlreadyStartedOptions {
                ignore_already_started_error: options.ignore_already_started,
            }),
            ..Default::default()
        };
        Ok(flow_start_options)
    }

    fn map_flow_config(&self, config: Option<&FlowConfig>) -> SdkResult<ProtoFlowConfig> {
        map_flow_config(config, self.options.worker_target_value())
    }

    fn call_empty<Response, Future, Call>(
        &self,
        operation: &'static str,
        flow_id: Option<&str>,
        requirement: FlowTargetRequirement,
        call: Call,
    ) -> SdkResult<()>
    where
        Call: FnOnce(FlowServiceClient<tonic::transport::Channel>) -> Future,
        Future: std::future::Future<Output = Result<tonic::Response<Response>, tonic::Status>>,
    {
        let service = self.service.clone();
        self.runtime
            .block_on(call(service))
            .map(|_| ())
            .map_err(|status| SdkError::from_status(status, operation, flow_id, requirement))
    }
}

pub(crate) fn map_flow_config(
    config: Option<&FlowConfig>,
    default_target: Option<&WorkerTarget>,
) -> SdkResult<ProtoFlowConfig> {
    let target = config
        .and_then(|config| config.worker_target.as_ref())
        .or(default_target);
    Ok(ProtoFlowConfig {
        active_step_search_mode: config
            .and_then(|config| config.active_step_search_mode)
            .map(|mode| match mode {
                ActiveStepSearchMode::All => ProtoSearchMode::EnabledForAll as i32,
                ActiveStepSearchMode::WithWaitFor => {
                    ProtoSearchMode::EnabledForStepsWithWaitFor as i32
                }
                ActiveStepSearchMode::Disabled => ProtoSearchMode::Disabled as i32,
            }),
        continue_as_new_threshold: config
            .and_then(|config| config.continue_as_new_threshold)
            .map(i32::try_from)
            .transpose()
            .map_err(|_| invalid("continue-as-new threshold exceeds int32"))?,
        continue_as_new_page_size_in_bytes: config
            .and_then(|config| config.continue_as_new_page_size_bytes)
            .map(i32::try_from)
            .transpose()
            .map_err(|_| invalid("continue-as-new page size exceeds int32"))?,
        step_durability: config
            .and_then(|config| config.step_durability)
            .map(|durability| {
                (match durability {
                    StepDurability::Default => ProtoStepDurability::Unspecified,
                    StepDurability::Sync => ProtoStepDurability::Sync,
                    StepDurability::Async => ProtoStepDurability::Async,
                }) as i32
            }),
        worker_target: target.map(map_worker_target),
        attribute_store_names: config.and_then(|config| {
            config
                .attribute_store_names
                .as_ref()
                .map(|names| AttributeStoreNames {
                    names: names.clone(),
                })
        }),
    })
}

pub(crate) fn map_flow_retry(retry: RetryPolicy) -> SdkResult<FlowRetryPolicy> {
    Ok(FlowRetryPolicy {
        initial_interval_seconds: optional_seconds(retry.initial_interval)?,
        backoff_coefficient: retry.backoff_coefficient.unwrap_or_default() as f32,
        maximum_interval_seconds: optional_seconds(retry.maximum_interval)?,
        maximum_attempts: i32::try_from(retry.maximum_attempts.unwrap_or_default())
            .map_err(|_| invalid("maximum attempts exceed int32"))?,
    })
}

fn map_worker_target(target: &WorkerTarget) -> ProtoWorkerTarget {
    ProtoWorkerTarget {
        address: target.address().to_string(),
        is_headless_address: target.is_headless(),
    }
}

pub(crate) fn map_flow_status(status: i32) -> SdkResult<FlowStatus> {
    match ProtoFlowStatus::try_from(status).ok() {
        Some(ProtoFlowStatus::Running) => Ok(FlowStatus::Running),
        Some(ProtoFlowStatus::Completed) => Ok(FlowStatus::Completed),
        Some(ProtoFlowStatus::Failed) => Ok(FlowStatus::Failed),
        Some(ProtoFlowStatus::ServerSideTimeoutInternalOnly) => {
            Ok(FlowStatus::ServerSideTimeoutInternalOnly)
        }
        Some(ProtoFlowStatus::Terminated) => Ok(FlowStatus::Terminated),
        Some(ProtoFlowStatus::Canceled) => Ok(FlowStatus::Canceled),
        Some(ProtoFlowStatus::ContinuedAsNew) => Ok(FlowStatus::ContinuedAsNew),
        _ => Err(invalid(format!("unknown Flow status {status}"))),
    }
}

pub(crate) fn map_flow_error_type(error_type: i32) -> Option<FlowErrorType> {
    use dex_protocol::dex::FlowErrorType as ProtoFlowErrorType;

    match ProtoFlowErrorType::try_from(error_type).ok() {
        Some(ProtoFlowErrorType::StepDecisionFailingFlow) => {
            Some(FlowErrorType::StepDecisionFailed)
        }
        Some(ProtoFlowErrorType::ClientApiFailingFlow) => Some(FlowErrorType::ClientApiFailed),
        Some(ProtoFlowErrorType::WorkerApiFail) => Some(FlowErrorType::WorkerApiFailed),
        Some(ProtoFlowErrorType::InvalidUserFlowCode) => Some(FlowErrorType::InvalidUserFlowCode),
        Some(ProtoFlowErrorType::Internal) => Some(FlowErrorType::Internal),
        Some(ProtoFlowErrorType::FlowTimeout) => Some(FlowErrorType::FlowTimeout),
        _ => None,
    }
}

pub(crate) fn map_flow_timeout_policy(
    flow: &crate::registry::RegisteredFlow,
    timeout_seconds: i32,
    policy: FlowTimeoutPolicy,
) -> SdkResult<i32> {
    if timeout_seconds == 0 {
        if policy != FlowTimeoutPolicy::Default {
            return Err(invalid("Flow timeout policy requires a positive timeout"));
        }
        return Ok(ProtoFlowTimeoutPolicy::Unspecified as i32);
    }
    let resolved = match policy {
        FlowTimeoutPolicy::Default if flow.handler.has_timeout_handler() => {
            FlowTimeoutPolicy::Handler
        }
        FlowTimeoutPolicy::Default => FlowTimeoutPolicy::Fail,
        configured => configured,
    };
    if resolved == FlowTimeoutPolicy::Handler && !flow.handler.has_timeout_handler() {
        return Err(invalid(format!(
            "Flow {} has no timeout handler",
            flow.name
        )));
    }
    Ok(match resolved {
        FlowTimeoutPolicy::Default => ProtoFlowTimeoutPolicy::Unspecified,
        FlowTimeoutPolicy::Fail => ProtoFlowTimeoutPolicy::Fail,
        FlowTimeoutPolicy::Cancel => ProtoFlowTimeoutPolicy::Cancel,
        FlowTimeoutPolicy::Handler => ProtoFlowTimeoutPolicy::Handler,
    } as i32)
}

fn timestamp(timestamp: prost_types::Timestamp) -> SystemTime {
    if timestamp.seconds >= 0 {
        SystemTime::UNIX_EPOCH
            + Duration::from_secs(timestamp.seconds as u64)
            + Duration::from_nanos(timestamp.nanos.max(0) as u64)
    } else {
        SystemTime::UNIX_EPOCH - Duration::from_secs(timestamp.seconds.unsigned_abs())
    }
}

fn rfc3339(time: SystemTime) -> SdkResult<String> {
    let duration = time
        .duration_since(SystemTime::UNIX_EPOCH)
        .map_err(service_error)?;
    let days = (duration.as_secs() / 86_400) as i64;
    let seconds = duration.as_secs() % 86_400;
    let (year, month, day) = civil_from_days(days);
    let hour = seconds / 3_600;
    let minute = seconds % 3_600 / 60;
    let second = seconds % 60;
    Ok(format!(
        "{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z"
    ))
}

fn civil_from_days(days: i64) -> (i64, i64, i64) {
    let days = days + 719_468;
    let era = if days >= 0 { days } else { days - 146_096 } / 146_097;
    let day_of_era = days - era * 146_097;
    let year_of_era =
        (day_of_era - day_of_era / 1_460 + day_of_era / 36_524 - day_of_era / 146_096) / 365;
    let mut year = year_of_era + era * 400;
    let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100);
    let month_prime = (5 * day_of_year + 2) / 153;
    let day = day_of_year - (153 * month_prime + 2) / 5 + 1;
    let month = month_prime + if month_prime < 10 { 3 } else { -9 };
    year += i64::from(month <= 2);
    (year, month, day)
}

fn optional_seconds(duration: Option<Duration>) -> SdkResult<i32> {
    duration
        .map(seconds32)
        .transpose()
        .map(Option::unwrap_or_default)
}

fn seconds32(duration: Duration) -> SdkResult<i32> {
    if duration.subsec_nanos() != 0 {
        return Err(invalid("Duration must use whole seconds"));
    }
    i32::try_from(duration.as_secs()).map_err(|_| invalid("Duration exceeds int32"))
}

fn endpoint_address(address: &str) -> String {
    if address.contains("://") {
        address.to_string()
    } else {
        format!("http://{address}")
    }
}

fn require_name(value: &str, kind: &str) -> SdkResult<()> {
    if value.is_empty() {
        Err(invalid(format!("{kind} is required")))
    } else {
        Ok(())
    }
}

fn sdk_handler_error(error: impl std::fmt::Display) -> SdkError {
    invalid(error.to_string())
}

fn invalid(message: impl Into<String>) -> SdkError {
    SdkError::InvalidArgument {
        message: message.into(),
    }
}

fn service_error(error: impl std::fmt::Display) -> SdkError {
    SdkError::Service {
        service: ServiceError::local("client", error.to_string()),
    }
}

#[cfg(test)]
mod tests {
    use super::map_flow_config;
    use crate::FlowConfig;

    #[test]
    fn attribute_store_names_preserve_protocol_presence() {
        let absent = map_flow_config(Some(&FlowConfig::new()), None).expect("map absent config");
        let named = map_flow_config(
            Some(
                &FlowConfig::new()
                    .attribute_store_names(vec!["profiles".to_owned(), "audit".to_owned()]),
            ),
            None,
        )
        .expect("map named config");
        let disabled =
            map_flow_config(Some(&FlowConfig::new().attribute_store_names(vec![])), None)
                .expect("map disabled config");

        assert_eq!(absent.attribute_store_names, None);
        assert_eq!(
            named.attribute_store_names.expect("named stores").names,
            vec!["profiles", "audit"]
        );
        assert_eq!(
            disabled
                .attribute_store_names
                .expect("disabled stores")
                .names,
            Vec::<String>::new()
        );
    }
}