dex_sdk/
context.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::collections::{HashMap, HashSet};
10use std::sync::atomic::{AtomicBool, Ordering};
11use std::sync::{Arc, Condvar, Mutex};
12use std::time::{Duration, SystemTime};
13
14use dex_protocol::dex::{
15    AttributeWrite, ChannelInfo, ChannelMessage, ConditionResults, ConditionStatus,
16    Context as ProtoContext, Kv, Value as ProtoValue, value,
17};
18
19use crate::persistence::PersistenceKind;
20use crate::registry::{RegisteredFlow, decode_instance, physical_name};
21use crate::value_mapper;
22use crate::{Attribute, AttributeMap, Channel, ChannelMap, HandlerError, HandlerResult, Value};
23
24#[derive(Clone, Copy, Eq, PartialEq)]
25pub(crate) enum InvocationMethod {
26    WaitFor,
27    Execute,
28    Rpc,
29}
30
31/// Provides invocation metadata and staged durable changes to Step and RPC handlers.
32///
33/// A Context belongs to one handler attempt and must not outlive the call. Attribute writes,
34/// Channel publications, locals, and events become visible atomically only when the handler returns
35/// successfully. Application values are freshly decoded and owned by the invocation.
36pub struct Context {
37    method: InvocationMethod,
38    flow: RegisteredFlow,
39    metadata: ProtoContext,
40    attributes: HashMap<String, ProtoValue>,
41    locals: HashMap<String, ProtoValue>,
42    condition_results: Option<ConditionResults>,
43    channel_infos: HashMap<String, ChannelInfo>,
44    attribute_writes: HashMap<String, AttributeWrite>,
45    local_writes: HashMap<String, Kv>,
46    events: Vec<Kv>,
47    event_names: HashSet<String>,
48    publications: Vec<ChannelMessage>,
49    cancellation: InvocationCancellation,
50}
51
52impl Context {
53    /// Returns the application-assigned Flow ID.
54    pub fn flow_id(&self) -> &str {
55        &self.metadata.flow_id
56    }
57
58    /// Returns the server-assigned run ID for this execution.
59    pub fn run_id(&self) -> &str {
60        &self.metadata.run_id
61    }
62
63    /// Returns when the current Flow run started.
64    pub fn flow_started_at(&self) -> SystemTime {
65        system_time(self.metadata.flow_started_timestamp)
66    }
67
68    /// Returns the current Step execution ID, or an empty string for an RPC.
69    pub fn step_execution_id(&self) -> &str {
70        &self.metadata.step_execution_id
71    }
72
73    /// Returns the predecessor Step execution ID, or an empty string when absent.
74    pub fn from_step_execution_id(&self) -> &str {
75        &self.metadata.from_step_execution_id
76    }
77
78    /// Returns when the first attempt of this handler invocation started.
79    pub fn first_attempt_at(&self) -> SystemTime {
80        system_time(self.metadata.first_attempt_timestamp)
81    }
82
83    /// Returns the one-based handler attempt number.
84    pub fn attempt(&self) -> u32 {
85        u32::try_from(self.metadata.attempt).unwrap_or_default()
86    }
87
88    /// Returns whether any timer condition completed for this `execute` invocation.
89    pub fn has_any_timer_fired(&self) -> bool {
90        self.condition_results.as_ref().is_some_and(|results| {
91            results
92                .timer_results
93                .iter()
94                .any(|timer| timer.condition_status == ConditionStatus::Completed as i32)
95        })
96    }
97
98    /// Returns whether the zero-based timer at `index` completed.
99    ///
100    /// Missing indexes return `false`.
101    pub fn has_timer_fired(&self, index: usize) -> bool {
102        self.condition_results
103            .as_ref()
104            .and_then(|results| results.timer_results.get(index))
105            .is_some_and(|timer| timer.condition_status == ConditionStatus::Completed as i32)
106    }
107
108    /// Returns whether `execute` is running because `wait_for` exhausted retries with Proceed.
109    pub fn wait_for_method_failed(&self) -> bool {
110        self.condition_results
111            .as_ref()
112            .is_some_and(|results| results.wait_for_failed)
113    }
114
115    /// Blocks the current thread until Dex cancels this handler attempt.
116    ///
117    /// Use this only for interruptible application work; ordinary handlers should return normally.
118    pub fn wait_for_cancellation(&self) {
119        self.cancellation.wait();
120    }
121
122    /// Returns immediately with the current handler-cancellation state.
123    pub fn is_cancelled(&self) -> bool {
124        self.cancellation.is_cancelled()
125    }
126
127    /// Stages a Step-execution-local value visible to later attempts of the same execution.
128    ///
129    /// # Errors
130    ///
131    /// Returns [`HandlerError`] for an empty key or an encoding failure.
132    pub fn set_step_execution_local<T: Value>(&mut self, key: &str, value: T) -> HandlerResult<()> {
133        require_name(key, "step-execution local key")?;
134        let value = value_mapper::encode_handler(&value)?;
135        self.local_writes.insert(
136            key.to_string(),
137            Kv {
138                key: key.to_string(),
139                value: Some(value),
140            },
141        );
142        Ok(())
143    }
144
145    /// Reads a required Step-execution-local value, including writes staged by this attempt.
146    ///
147    /// # Errors
148    ///
149    /// Returns [`HandlerError`] for an empty key, a missing value, or a decoding failure.
150    pub fn step_execution_local<T: Value>(&self, key: &str) -> HandlerResult<T> {
151        require_name(key, "step-execution local key")?;
152        let value = self
153            .local_writes
154            .get(key)
155            .and_then(|entry| entry.value.as_ref())
156            .or_else(|| self.locals.get(key))
157            .ok_or_else(|| HandlerError::new(format!("step-execution local {key} is missing")))?;
158        value_mapper::decode_handler(value)
159    }
160
161    /// Records one named diagnostic event with a typed payload.
162    ///
163    /// Event names must be unique within the invocation.
164    ///
165    /// # Errors
166    ///
167    /// Returns [`HandlerError`] for an empty or duplicate name or an encoding failure.
168    pub fn record_event<T: Value>(&mut self, name: &str, value: T) -> HandlerResult<()> {
169        require_name(name, "event name")?;
170        if !self.event_names.insert(name.to_string()) {
171            return Err(HandlerError::new(format!(
172                "event was already recorded: {name}"
173            )));
174        }
175        self.events.push(Kv {
176            key: name.to_string(),
177            value: Some(value_mapper::encode_handler(&value)?),
178        });
179        Ok(())
180    }
181
182    /// Reads an Attribute value, including writes staged by this invocation.
183    pub fn get_attribute<T: Value>(&self, attribute: &Attribute<T>) -> HandlerResult<Option<T>> {
184        self.get_attribute_value(attribute.name(), PersistenceKind::Attribute, None)
185    }
186
187    /// Reads one Attribute-map instance, including staged writes.
188    pub fn get_attribute_map<T: Value>(
189        &self,
190        attribute: &AttributeMap<T>,
191        instance: &str,
192    ) -> HandlerResult<Option<T>> {
193        self.get_attribute_value(
194            attribute.name(),
195            PersistenceKind::AttributeMap,
196            Some(instance),
197        )
198    }
199
200    /// Stages a typed Attribute write.
201    pub fn set_attribute<T: Value>(
202        &mut self,
203        attribute: &Attribute<T>,
204        value: T,
205    ) -> HandlerResult<()> {
206        self.set_attribute_value(
207            attribute.name(),
208            PersistenceKind::Attribute,
209            None,
210            value_mapper::encode_handler(&value)?,
211            attribute.index().map(|index| index.proto_config(false)),
212            attribute.sync_config(),
213        )
214    }
215
216    /// Stages a typed write for one Attribute-map instance.
217    pub fn set_attribute_map<T: Value>(
218        &mut self,
219        attribute: &AttributeMap<T>,
220        instance: &str,
221        value: T,
222    ) -> HandlerResult<()> {
223        self.set_attribute_value(
224            attribute.name(),
225            PersistenceKind::AttributeMap,
226            Some(instance),
227            value_mapper::encode_handler(&value)?,
228            attribute.index().map(|index| index.proto_config(true)),
229            attribute.sync_config(),
230        )
231    }
232
233    /// Stages deletion of an Attribute value.
234    pub fn delete_attribute<T>(&mut self, attribute: &Attribute<T>) -> HandlerResult<()> {
235        self.set_attribute_value(
236            attribute.name(),
237            PersistenceKind::Attribute,
238            None,
239            value_mapper::deletion(),
240            attribute.index().map(|index| index.proto_config(false)),
241            attribute.sync_config(),
242        )
243    }
244
245    /// Stages deletion of one Attribute-map instance.
246    pub fn delete_attribute_map<T>(
247        &mut self,
248        attribute: &AttributeMap<T>,
249        instance: &str,
250    ) -> HandlerResult<()> {
251        self.set_attribute_value(
252            attribute.name(),
253            PersistenceKind::AttributeMap,
254            Some(instance),
255            value_mapper::deletion(),
256            attribute.index().map(|index| index.proto_config(true)),
257            attribute.sync_config(),
258        )
259    }
260
261    pub(crate) fn attribute_map_keys<T>(
262        &self,
263        attribute: &AttributeMap<T>,
264    ) -> HandlerResult<Vec<String>> {
265        self.registered_name(
266            attribute.name(),
267            PersistenceKind::AttributeMap,
268            Some("key-check"),
269        )?;
270        let prefix = format!("{}/", attribute.name());
271        let mut physical_keys = self
272            .attributes
273            .keys()
274            .filter(|key| key.starts_with(&prefix))
275            .cloned()
276            .collect::<HashSet<_>>();
277        for (key, write) in &self.attribute_writes {
278            if !key.starts_with(&prefix) {
279                continue;
280            }
281            if matches!(
282                write.value.as_ref().and_then(|value| value.kind.as_ref()),
283                Some(value::Kind::NullValue(_))
284            ) {
285                physical_keys.remove(key);
286            } else {
287                physical_keys.insert(key.clone());
288            }
289        }
290        sorted_instance_keys(&prefix, physical_keys)
291    }
292
293    /// Stages one typed Channel publication.
294    pub fn publish<T: Value>(&mut self, channel: &Channel<T>, value: T) -> HandlerResult<()> {
295        self.publish_value(
296            channel.name(),
297            PersistenceKind::Channel,
298            None,
299            value_mapper::encode_handler(&value)?,
300        )
301    }
302
303    /// Stages one typed publication to a Channel-map instance.
304    pub fn publish_map<T: Value>(
305        &mut self,
306        channel: &ChannelMap<T>,
307        instance: &str,
308        value: T,
309    ) -> HandlerResult<()> {
310        self.publish_value(
311            channel.name(),
312            PersistenceKind::ChannelMap,
313            Some(instance),
314            value_mapper::encode_handler(&value)?,
315        )
316    }
317
318    /// Returns the current invocation snapshot's Channel queue size.
319    pub fn channel_size<T>(&self, channel: &Channel<T>) -> HandlerResult<usize> {
320        self.channel_size_value(channel.name(), PersistenceKind::Channel, None)
321    }
322
323    /// Returns one Channel-map instance's queue size.
324    pub fn channel_map_size<T>(
325        &self,
326        channel: &ChannelMap<T>,
327        instance: &str,
328    ) -> HandlerResult<usize> {
329        self.channel_size_value(channel.name(), PersistenceKind::ChannelMap, Some(instance))
330    }
331
332    pub(crate) fn channel_map_keys<T>(
333        &self,
334        channel: &ChannelMap<T>,
335    ) -> HandlerResult<Vec<String>> {
336        if self.method != InvocationMethod::Rpc {
337            return Err(HandlerError::new(
338                "ChannelMap introspection requires an RPC invocation",
339            ));
340        }
341        self.registered_name(
342            channel.name(),
343            PersistenceKind::ChannelMap,
344            Some("key-check"),
345        )?;
346        let prefix = format!("{}/", channel.name());
347        let physical_keys = self
348            .channel_infos
349            .iter()
350            .filter(|(key, info)| key.starts_with(&prefix) && info.size > 0)
351            .map(|(key, _)| key.clone())
352            .collect::<HashSet<_>>();
353        sorted_instance_keys(&prefix, physical_keys)
354    }
355
356    /// Decodes the messages consumed by a satisfied Channel condition.
357    pub fn channel_results<T: Value>(&self, channel: &Channel<T>) -> HandlerResult<Vec<T>> {
358        self.channel_results_value(channel.name(), PersistenceKind::Channel, None)
359    }
360
361    /// Decodes messages consumed by one satisfied Channel-map condition.
362    pub fn channel_map_results<T: Value>(
363        &self,
364        channel: &ChannelMap<T>,
365        instance: &str,
366    ) -> HandlerResult<Vec<T>> {
367        self.channel_results_value(channel.name(), PersistenceKind::ChannelMap, Some(instance))
368    }
369
370    pub(crate) fn new(
371        method: InvocationMethod,
372        flow: RegisteredFlow,
373        metadata: ProtoContext,
374        attributes: Vec<Kv>,
375        locals: Vec<Kv>,
376        condition_results: Option<ConditionResults>,
377        channel_infos: HashMap<String, ChannelInfo>,
378    ) -> HandlerResult<Self> {
379        Ok(Self {
380            method,
381            flow,
382            metadata,
383            attributes: map_values("Attribute", attributes)?,
384            locals: map_values("step-execution local", locals)?,
385            condition_results,
386            channel_infos,
387            attribute_writes: HashMap::new(),
388            local_writes: HashMap::new(),
389            events: Vec::new(),
390            event_names: HashSet::new(),
391            publications: Vec::new(),
392            cancellation: InvocationCancellation::new(),
393        })
394    }
395
396    pub(crate) fn cancellation(&self) -> InvocationCancellation {
397        self.cancellation.clone()
398    }
399
400    pub(crate) fn take_outputs(
401        self,
402    ) -> (Vec<AttributeWrite>, Vec<Kv>, Vec<Kv>, Vec<ChannelMessage>) {
403        (
404            self.attribute_writes.into_values().collect(),
405            self.local_writes.into_values().collect(),
406            self.events,
407            self.publications,
408        )
409    }
410
411    fn get_attribute_value<T: Value>(
412        &self,
413        name: &str,
414        kind: PersistenceKind,
415        instance: Option<&str>,
416    ) -> HandlerResult<Option<T>> {
417        let key = self.registered_name(name, kind, instance)?;
418        if let Some(write) = self.attribute_writes.get(&key) {
419            let value = write
420                .value
421                .as_ref()
422                .ok_or_else(|| HandlerError::new(format!("Attribute write {key} has no Value")))?;
423            if matches!(value.kind, Some(value::Kind::NullValue(_))) {
424                return Ok(None);
425            }
426            return value_mapper::decode_handler(value).map(Some);
427        }
428        self.attributes
429            .get(&key)
430            .map(value_mapper::decode_handler)
431            .transpose()
432    }
433
434    fn set_attribute_value(
435        &mut self,
436        name: &str,
437        kind: PersistenceKind,
438        instance: Option<&str>,
439        value: ProtoValue,
440        index_config: Option<dex_protocol::dex::IndexConfig>,
441        sync_config: Option<dex_protocol::dex::AttributeSyncConfig>,
442    ) -> HandlerResult<()> {
443        let key = self.registered_name(name, kind, instance)?;
444        self.attribute_writes.insert(
445            key.clone(),
446            AttributeWrite {
447                key,
448                value: Some(value),
449                index_config,
450                sync_config,
451            },
452        );
453        Ok(())
454    }
455
456    fn publish_value(
457        &mut self,
458        name: &str,
459        kind: PersistenceKind,
460        instance: Option<&str>,
461        value: ProtoValue,
462    ) -> HandlerResult<()> {
463        let channel_name = self.registered_name(name, kind, instance)?;
464        self.publications.push(ChannelMessage {
465            channel_name: channel_name.clone(),
466            value: Some(value),
467        });
468        if self.method == InvocationMethod::Rpc {
469            self.channel_infos
470                .entry(channel_name)
471                .and_modify(|info| info.size += 1)
472                .or_insert(ChannelInfo { size: 1 });
473        }
474        Ok(())
475    }
476
477    fn channel_size_value(
478        &self,
479        name: &str,
480        kind: PersistenceKind,
481        instance: Option<&str>,
482    ) -> HandlerResult<usize> {
483        let channel_name = self.registered_name(name, kind, instance)?;
484        Ok(self
485            .channel_infos
486            .get(&channel_name)
487            .map_or(0, |info| usize::try_from(info.size).unwrap_or_default()))
488    }
489
490    fn channel_results_value<T: Value>(
491        &self,
492        name: &str,
493        kind: PersistenceKind,
494        instance: Option<&str>,
495    ) -> HandlerResult<Vec<T>> {
496        let channel_name = self.registered_name(name, kind, instance)?;
497        let Some(results) = &self.condition_results else {
498            return Ok(Vec::new());
499        };
500        let mut decoded = Vec::new();
501        for result in &results.channel_results {
502            if result.channel_name == channel_name
503                && result.condition_status == ConditionStatus::Completed as i32
504            {
505                for value in &result.values {
506                    decoded.push(value_mapper::decode_handler(value)?);
507                }
508            }
509        }
510        Ok(decoded)
511    }
512
513    fn registered_name(
514        &self,
515        name: &str,
516        kind: PersistenceKind,
517        instance: Option<&str>,
518    ) -> HandlerResult<String> {
519        let definition = self.flow.persistence.get(name).ok_or_else(|| {
520            HandlerError::new(format!(
521                "persistence definition does not belong to Flow: {name}"
522            ))
523        })?;
524        if definition.kind != kind {
525            return Err(HandlerError::new(format!(
526                "persistence definition kind does not match: {name}"
527            )));
528        }
529        match kind {
530            PersistenceKind::AttributeMap | PersistenceKind::ChannelMap => instance
531                .map(|instance| physical_name(name, instance))
532                .ok_or_else(|| HandlerError::new(format!("{name} requires an instance"))),
533            PersistenceKind::Attribute | PersistenceKind::Channel if instance.is_none() => {
534                Ok(name.to_string())
535            }
536            _ => Err(HandlerError::new(format!(
537                "static persistence definition cannot use an instance: {name}"
538            ))),
539        }
540    }
541}
542
543fn sorted_instance_keys(
544    prefix: &str,
545    physical_keys: HashSet<String>,
546) -> HandlerResult<Vec<String>> {
547    let mut keys = physical_keys
548        .iter()
549        .map(|key| decode_instance(key, prefix).map_err(HandlerError::new))
550        .collect::<HandlerResult<Vec<_>>>()?;
551    keys.sort();
552    Ok(keys)
553}
554
555#[derive(Clone)]
556pub(crate) struct InvocationCancellation {
557    inner: Arc<InvocationCancellationInner>,
558}
559
560struct InvocationCancellationInner {
561    cancelled: AtomicBool,
562    mutex: Mutex<()>,
563    condition: Condvar,
564}
565
566impl InvocationCancellation {
567    fn new() -> Self {
568        Self {
569            inner: Arc::new(InvocationCancellationInner {
570                cancelled: AtomicBool::new(false),
571                mutex: Mutex::new(()),
572                condition: Condvar::new(),
573            }),
574        }
575    }
576
577    fn wait(&self) {
578        let mut guard = self.inner.mutex.lock().expect("cancellation lock");
579        while !self.is_cancelled() {
580            guard = self.inner.condition.wait(guard).expect("cancellation lock");
581        }
582    }
583
584    fn is_cancelled(&self) -> bool {
585        self.inner.cancelled.load(Ordering::Acquire)
586    }
587
588    pub(crate) fn cancel(&self) {
589        let _guard = self.inner.mutex.lock().expect("cancellation lock");
590        self.inner.cancelled.store(true, Ordering::Release);
591        self.inner.condition.notify_all();
592    }
593}
594
595fn system_time(timestamp: i64) -> SystemTime {
596    if timestamp >= 0 {
597        SystemTime::UNIX_EPOCH + Duration::from_secs(timestamp as u64)
598    } else {
599        SystemTime::UNIX_EPOCH - Duration::from_secs(timestamp.unsigned_abs())
600    }
601}
602
603fn map_values(kind: &str, entries: Vec<Kv>) -> HandlerResult<HashMap<String, ProtoValue>> {
604    let mut mapped = HashMap::new();
605    for entry in entries {
606        require_name(&entry.key, kind)?;
607        let value = entry
608            .value
609            .ok_or_else(|| HandlerError::new(format!("{kind} {} has no Value", entry.key)))?;
610        if mapped.insert(entry.key.clone(), value).is_some() {
611            return Err(HandlerError::new(format!("duplicate {kind} {}", entry.key)));
612        }
613    }
614    Ok(mapped)
615}
616
617fn require_name(value: &str, kind: &str) -> HandlerResult<()> {
618    if value.is_empty() {
619        Err(HandlerError::new(format!("{kind} is required")))
620    } else {
621        Ok(())
622    }
623}
624
625#[cfg(test)]
626mod tests {
627    use super::{Context, InvocationMethod};
628    use crate::{
629        AttributeMap, ChannelMap, Flow, PersistenceSchema, Registry, registry::physical_name,
630        value_mapper,
631    };
632    use dex_protocol::dex::{ChannelInfo, Context as ProtoContext, Kv};
633    use std::collections::HashMap;
634
635    struct MapFlow {
636        attributes: AttributeMap<String>,
637        channels: ChannelMap<String>,
638    }
639
640    impl Flow for MapFlow {
641        type StartInput = ();
642
643        fn persistence(&self) -> PersistenceSchema {
644            PersistenceSchema::new()
645                .attribute_map(&self.attributes)
646                .channel_map(&self.channels)
647        }
648    }
649
650    #[test]
651    fn map_introspection_tracks_buffered_changes() {
652        let flow = MapFlow {
653            attributes: AttributeMap::new("items"),
654            channels: ChannelMap::new("messages"),
655        };
656        let registry = Registry::new().register(flow).expect("register map Flow");
657        let registered = registry.flow("MapFlow").expect("lookup map Flow").clone();
658        let attributes = AttributeMap::<String>::new("items");
659        let channels = ChannelMap::<String>::new("messages");
660        let special = "special / key";
661        let mut context = Context::new(
662            InvocationMethod::Rpc,
663            registered,
664            ProtoContext::default(),
665            vec![
666                Kv {
667                    key: physical_name("items", special),
668                    value: Some(value_mapper::encode_handler(&"initial".to_string()).unwrap()),
669                },
670                Kv {
671                    key: physical_name("items", "z"),
672                    value: Some(value_mapper::encode_handler(&"remove".to_string()).unwrap()),
673                },
674            ],
675            Vec::new(),
676            None,
677            HashMap::from([
678                (physical_name("messages", special), ChannelInfo { size: 1 }),
679                (physical_name("messages", "empty"), ChannelInfo { size: 0 }),
680            ]),
681        )
682        .expect("create RPC Context");
683
684        assert_eq!(
685            vec![special.to_string(), "z".to_string()],
686            attributes.all_instance_keys(&context).unwrap()
687        );
688        attributes
689            .set(&mut context, "a", "added".to_string())
690            .unwrap();
691        attributes.delete(&mut context, "z").unwrap();
692        assert_eq!(
693            vec!["a".to_string(), special.to_string()],
694            attributes.all_instance_keys(&context).unwrap()
695        );
696        assert_eq!(2, attributes.map_size(&context).unwrap());
697
698        assert_eq!(
699            vec![special.to_string()],
700            channels.all_instance_keys(&context).unwrap()
701        );
702        channels
703            .publish(&mut context, "a", "published".to_string())
704            .unwrap();
705        assert_eq!(
706            vec!["a".to_string(), special.to_string()],
707            channels.all_instance_keys(&context).unwrap()
708        );
709        assert_eq!(2, channels.map_size(&context).unwrap());
710    }
711}
dex_sdk/context.rs

dex_sdk/
context.rs
