aion_worker/
activity.rs

1//! `Activity` trait, `ActivityFailure`, and typed registration.
2
3use std::any::Any;
4use std::collections::{BTreeMap, BTreeSet};
5use std::future::Future;
6use std::panic::AssertUnwindSafe;
7use std::pin::Pin;
8
9use aion_core::{ActivityError, ActivityErrorKind, Payload};
10use async_trait::async_trait;
11use futures::FutureExt;
12use serde::Serialize;
13use serde::de::DeserializeOwned;
14use tracing::error;
15
16use crate::context::ActivityContext;
17use crate::error::{MissingActivityHandler, WorkerError};
18use crate::protocol::ActivityTask;
19use crate::runtime::loop_::{ActivityDispatcher, DispatchOutcome};
20
21/// Explicit retryability classification for an activity failure.
22#[derive(Clone, Debug, PartialEq, Eq)]
23pub enum Classification {
24    /// The engine may retry the activity according to policy.
25    Retryable,
26    /// The provider refused under policy; route before ordinary retry handling.
27    PolicyRefused,
28    /// The activity failure is permanent and must not be retried.
29    Terminal,
30}
31
32/// Handler-returned failure with explicit retryability classification.
33#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
34#[error("{message}")]
35pub struct ActivityFailure {
36    classification: Classification,
37    message: String,
38    detail: Option<Payload>,
39}
40
41impl ActivityFailure {
42    /// Creates a retryable activity failure.
43    #[must_use]
44    pub fn retryable(message: impl Into<String>) -> Self {
45        Self::new(Classification::Retryable, message, None)
46    }
47
48    /// Creates a provider-policy-refused activity failure.
49    #[must_use]
50    pub fn policy_refused(message: impl Into<String>) -> Self {
51        Self::new(Classification::PolicyRefused, message, None)
52    }
53
54    /// Creates a terminal activity failure.
55    #[must_use]
56    pub fn terminal(message: impl Into<String>) -> Self {
57        Self::new(Classification::Terminal, message, None)
58    }
59
60    /// Attaches opaque structured detail to this failure.
61    #[must_use]
62    pub fn with_detail(mut self, detail: Payload) -> Self {
63        self.detail = Some(detail);
64        self
65    }
66
67    /// Returns the explicit retryability classification.
68    #[must_use]
69    pub const fn classification(&self) -> &Classification {
70        &self.classification
71    }
72
73    /// Returns the human-readable failure message.
74    #[must_use]
75    pub fn message(&self) -> &str {
76        &self.message
77    }
78
79    /// Returns the optional structured failure detail.
80    #[must_use]
81    pub const fn detail(&self) -> Option<&Payload> {
82        self.detail.as_ref()
83    }
84
85    fn new(
86        classification: Classification,
87        message: impl Into<String>,
88        detail: Option<Payload>,
89    ) -> Self {
90        Self {
91            classification,
92            message: message.into(),
93            detail,
94        }
95    }
96}
97
98impl From<Classification> for ActivityErrorKind {
99    fn from(value: Classification) -> Self {
100        match value {
101            Classification::Retryable => Self::Retryable,
102            Classification::PolicyRefused => Self::PolicyRefused,
103            Classification::Terminal => Self::Terminal,
104        }
105    }
106}
107
108impl From<ActivityFailure> for ActivityError {
109    fn from(value: ActivityFailure) -> Self {
110        Self {
111            kind: ActivityErrorKind::from(value.classification),
112            message: value.message,
113            details: value.detail,
114        }
115    }
116}
117
118/// Boxed future returned by a typed activity handler.
119pub type HandlerFuture<'context, Output> =
120    Pin<Box<dyn Future<Output = Result<Output, ActivityFailure>> + Send + 'context>>;
121
122type BoxedHandler<Input, Output> = Box<
123    dyn for<'context> Fn(Input, &'context ActivityContext) -> HandlerFuture<'context, Output>
124        + Send
125        + Sync,
126>;
127
128/// Registry of typed activity handlers keyed by activity-type name.
129#[derive(Default)]
130pub struct ActivityRegistry {
131    handlers: BTreeMap<String, Box<dyn ErasedActivityHandler>>,
132    descriptors: BTreeMap<String, aion_package::ActivityDescriptor>,
133}
134
135impl ActivityRegistry {
136    /// Creates an empty activity registry.
137    #[must_use]
138    pub fn new() -> Self {
139        Self::default()
140    }
141
142    /// Registers one typed activity handler under an activity-type name.
143    ///
144    /// # Errors
145    ///
146    /// Returns [`WorkerError::Registration`] when the name is already registered.
147    pub fn register_activity<Input, Output, Handler>(
148        mut self,
149        activity_type: impl Into<String>,
150        handler: Handler,
151    ) -> Result<Self, WorkerError>
152    where
153        Input: Serialize + DeserializeOwned + Send + Sync + 'static,
154        Output: Serialize + Send + Sync + 'static,
155        Handler: for<'context> Fn(Input, &'context ActivityContext) -> HandlerFuture<'context, Output>
156            + Send
157            + Sync
158            + 'static,
159    {
160        let activity_type = activity_type.into();
161        if self.handlers.contains_key(&activity_type) {
162            return Err(WorkerError::registration(DuplicateActivityType {
163                activity_type,
164            }));
165        }
166        self.handlers
167            .insert(activity_type, Box::new(TypedHandler::new(handler)));
168        Ok(self)
169    }
170
171    /// Registers a typed handler and mechanically derives its wire contract
172    /// from the same concrete serde types used at dispatch.
173    ///
174    /// # Errors
175    ///
176    /// Returns [`WorkerError::Registration`] for a duplicate name or
177    /// [`WorkerError::Encode`] if a generated schema cannot be represented as JSON.
178    pub fn register_activity_with_contract<Input, Output, Handler>(
179        mut self,
180        activity_type: impl Into<String>,
181        handler: Handler,
182    ) -> Result<Self, WorkerError>
183    where
184        Input: Serialize + DeserializeOwned + schemars::JsonSchema + Send + Sync + 'static,
185        Output: Serialize + schemars::JsonSchema + Send + Sync + 'static,
186        Handler: for<'context> Fn(Input, &'context ActivityContext) -> HandlerFuture<'context, Output>
187            + Send
188            + Sync
189            + 'static,
190    {
191        let activity_type = activity_type.into();
192        let descriptor = activity_descriptor::<Input, Output>(activity_type.clone())?;
193        self = self.register_activity(activity_type.clone(), handler)?;
194        self.descriptors.insert(activity_type, descriptor);
195        Ok(self)
196    }
197
198    /// Registers a typed handler together with an EXPLICIT wire descriptor.
199    ///
200    /// For an activity whose schemas are owned by a declaration rather than by
201    /// Rust types, [`Self::register_activity_with_contract`] cannot help: there
202    /// is no `schemars` type to derive from, because the types live in the
203    /// `.awl` document. The handler is therefore untyped (`serde_json::Value`
204    /// in and out) while the advertisement is exact, supplied by whoever read
205    /// the declaration.
206    ///
207    /// Registering a handler with no descriptor is what makes a worker invisible
208    /// to contract admission — it advertises nothing, and a queue carrying any
209    /// deployed contract refuses it. This is the path that keeps a
210    /// declaration-driven worker admissible.
211    ///
212    /// # Errors
213    ///
214    /// Returns [`WorkerError::Registration`] when the name is already
215    /// registered, or when `descriptor` names a different activity than
216    /// `activity_type` — a descriptor advertised under the wrong name would
217    /// promise one action's schema for another's handler.
218    pub fn register_activity_with_descriptor<Input, Output, Handler>(
219        mut self,
220        activity_type: impl Into<String>,
221        descriptor: aion_package::ActivityDescriptor,
222        handler: Handler,
223    ) -> Result<Self, WorkerError>
224    where
225        Input: Serialize + DeserializeOwned + Send + Sync + 'static,
226        Output: Serialize + Send + Sync + 'static,
227        Handler: for<'context> Fn(Input, &'context ActivityContext) -> HandlerFuture<'context, Output>
228            + Send
229            + Sync
230            + 'static,
231    {
232        let activity_type = activity_type.into();
233        if descriptor.name != activity_type {
234            return Err(WorkerError::registration(DescriptorNameMismatch {
235                activity_type,
236                descriptor: descriptor.name,
237            }));
238        }
239        self = self.register_activity(activity_type.clone(), handler)?;
240        self.descriptors.insert(activity_type, descriptor);
241        Ok(self)
242    }
243
244    /// Returns true when no activity handlers have been registered.
245    #[must_use]
246    pub fn is_empty(&self) -> bool {
247        self.handlers.is_empty()
248    }
249
250    /// Returns the registered activity-type names in deterministic order.
251    #[must_use]
252    pub fn activity_types(&self) -> BTreeSet<String> {
253        self.handlers.keys().cloned().collect()
254    }
255
256    /// Returns committed typed activity descriptors in deterministic name order.
257    #[must_use]
258    pub fn activity_descriptors(&self) -> Vec<aion_package::ActivityDescriptor> {
259        self.descriptors.values().cloned().collect()
260    }
261}
262
263/// Mechanically derives a wire descriptor from concrete serde activity types.
264///
265/// # Errors
266///
267/// Returns [`WorkerError::Encode`] if a generated schema cannot be represented
268/// as JSON.
269pub fn activity_descriptor<Input, Output>(
270    name: impl Into<String>,
271) -> Result<aion_package::ActivityDescriptor, WorkerError>
272where
273    Input: schemars::JsonSchema,
274    Output: schemars::JsonSchema,
275{
276    let schema = || {
277        schemars::generate::SchemaSettings::draft2020_12()
278            .with(|settings| {
279                // AWL contracts carry their dialect at the document root; the
280                // registration envelope already fixes that dialect, so draft
281                // metadata and presentation-only `$defs` must not create drift.
282                settings.meta_schema = None;
283                settings.inline_subschemas = true;
284            })
285            .into_generator()
286    };
287    Ok(aion_package::ActivityDescriptor {
288        name: name.into(),
289        input_schema: serde_json::to_value(schema().into_root_schema_for::<Input>())
290            .map_err(WorkerError::encode)?,
291        output_schema: serde_json::to_value(schema().into_root_schema_for::<Output>())
292            .map_err(WorkerError::encode)?,
293    })
294}
295
296#[async_trait]
297impl ActivityDispatcher for ActivityRegistry {
298    async fn dispatch(
299        &self,
300        task: ActivityTask,
301        context: ActivityContext,
302    ) -> Result<DispatchOutcome, WorkerError> {
303        let Some(handler) = self.handlers.get(&task.activity_type) else {
304            return Err(WorkerError::registration(MissingActivityHandler {
305                activity_type: task.activity_type,
306            }));
307        };
308        handler.dispatch(task, context).await
309    }
310
311    fn activity_types(&self) -> BTreeSet<String> {
312        self.activity_types()
313    }
314}
315
316/// Backwards-compatible name for the typed activity registry used by the runtime.
317pub type TypedActivityDispatcher = ActivityRegistry;
318
319/// Decodes a payload into a typed value using the payload content-type tag.
320///
321/// # Errors
322///
323/// Returns [`WorkerError::Decode`] when the payload tag or bytes cannot produce
324/// the requested type.
325pub fn decode_payload<T>(payload: &Payload) -> Result<T, WorkerError>
326where
327    T: DeserializeOwned,
328{
329    let value = payload.to_json().map_err(WorkerError::decode)?;
330    serde_json::from_value(value).map_err(WorkerError::decode)
331}
332
333/// Encodes a typed value into the baseline JSON payload codec.
334///
335/// # Errors
336///
337/// Returns [`WorkerError::Encode`] when the value cannot be serialized.
338pub fn encode_payload<T>(value: &T) -> Result<Payload, WorkerError>
339where
340    T: Serialize,
341{
342    let value = serde_json::to_value(value).map_err(WorkerError::encode)?;
343    Payload::from_json(&value).map_err(WorkerError::encode)
344}
345
346#[async_trait]
347trait ErasedActivityHandler: Send + Sync {
348    async fn dispatch(
349        &self,
350        task: ActivityTask,
351        context: ActivityContext,
352    ) -> Result<DispatchOutcome, WorkerError>;
353}
354
355struct TypedHandler<Input, Output> {
356    handler: BoxedHandler<Input, Output>,
357}
358
359impl<Input, Output> TypedHandler<Input, Output> {
360    fn new(
361        handler: impl for<'context> Fn(
362            Input,
363            &'context ActivityContext,
364        ) -> HandlerFuture<'context, Output>
365        + Send
366        + Sync
367        + 'static,
368    ) -> Self {
369        Self {
370            handler: Box::new(handler),
371        }
372    }
373}
374
375#[async_trait]
376impl<Input, Output> ErasedActivityHandler for TypedHandler<Input, Output>
377where
378    Input: DeserializeOwned + Send + Sync + 'static,
379    Output: Serialize + Send + Sync + 'static,
380{
381    async fn dispatch(
382        &self,
383        task: ActivityTask,
384        context: ActivityContext,
385    ) -> Result<DispatchOutcome, WorkerError> {
386        let input = match decode_payload::<Input>(&task.input) {
387            Ok(input) => input,
388            Err(error) => {
389                error!(
390                    activity_type = %task.activity_type,
391                    activity_id = task.activity_id.sequence_position(),
392                    attempt = task.attempt,
393                    error = %error,
394                    "failed to decode activity input; reporting terminal activity failure"
395                );
396                let failure =
397                    ActivityFailure::terminal(format!("failed to decode activity input: {error}"));
398                return Ok(DispatchOutcome::Failed {
399                    failure: ActivityError::from(failure),
400                });
401            }
402        };
403        let handler_future =
404            match std::panic::catch_unwind(AssertUnwindSafe(|| (self.handler)(input, &context))) {
405                Ok(handler_future) => handler_future,
406                Err(panic) => return Ok(panic_failure(&task, &panic)),
407            };
408        let handler_result = AssertUnwindSafe(handler_future).catch_unwind().await;
409        let outcome = match handler_result {
410            Ok(Ok(output)) => DispatchOutcome::Completed {
411                output: encode_payload(&output)?,
412            },
413            Ok(Err(failure)) => DispatchOutcome::Failed {
414                failure: ActivityError::from(failure),
415            },
416            Err(panic) => panic_failure(&task, &panic),
417        };
418        Ok(outcome)
419    }
420}
421
422fn panic_failure(task: &ActivityTask, panic: &Box<dyn Any + Send>) -> DispatchOutcome {
423    let message = panic_message(panic);
424    error!(
425        activity_type = %task.activity_type,
426        activity_id = task.activity_id.sequence_position(),
427        attempt = task.attempt,
428        panic = %message,
429        "activity handler panicked; reporting retryable activity failure"
430    );
431    DispatchOutcome::Failed {
432        failure: ActivityError::from(ActivityFailure::retryable(format!(
433            "activity handler panicked: {message}"
434        ))),
435    }
436}
437
438fn panic_message(panic: &Box<dyn Any + Send>) -> String {
439    if let Some(message) = panic.downcast_ref::<&str>() {
440        return (*message).to_owned();
441    }
442    if let Some(message) = panic.downcast_ref::<String>() {
443        return message.clone();
444    }
445    String::from("unknown panic payload")
446}
447
448/// Error returned when an activity type is registered more than once.
449#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
450#[error("activity type `{activity_type}` already has a registered handler")]
451pub struct DuplicateActivityType {
452    /// Duplicate activity type name.
453    pub activity_type: String,
454}
455
456/// Error returned when an explicit descriptor names a different activity than
457/// the one it is registered under.
458///
459/// Advertising a schema under the wrong name promises one action's shape for
460/// another action's handler, which contract admission would then accept — so
461/// the mismatch is refused at registration instead.
462#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
463#[error(
464    "activity type `{activity_type}` was registered with a descriptor naming `{descriptor}`; a descriptor must name the activity it describes"
465)]
466pub struct DescriptorNameMismatch {
467    /// The name the handler is registered under.
468    pub activity_type: String,
469    /// The name the supplied descriptor carries.
470    pub descriptor: String,
471}
472
473#[cfg(test)]
474mod tests {
475    use aion_core::{ActivityError, ActivityId, ContentType, WorkflowId};
476    use aion_proto::{
477        ProtoActivityError, ProtoActivityErrorKind, ProtoActivityId, ProtoActivityTask,
478        ProtoPayload, ProtoWorkflowId,
479    };
480    use serde::{Deserialize, Serialize};
481
482    use super::{ActivityFailure, ActivityRegistry, decode_payload, encode_payload};
483    use crate::WorkerError;
484    use crate::runtime::{ActivityDispatcher, DispatchOutcome};
485
486    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
487    struct TestInput {
488        value: i32,
489    }
490
491    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
492    struct TestOutput {
493        doubled: i32,
494    }
495
496    #[test]
497    fn retryable_and_terminal_failures_map_to_distinct_wire_classifications() {
498        let retryable = ActivityFailure::retryable("temporary outage");
499        let terminal = ActivityFailure::terminal("invalid request");
500
501        let retryable_core = ActivityError::from(retryable);
502        let terminal_core = ActivityError::from(terminal);
503        let retryable_wire = ProtoActivityError::from(retryable_core);
504        let terminal_wire = ProtoActivityError::from(terminal_core);
505
506        assert_eq!(
507            retryable_wire.kind,
508            ProtoActivityErrorKind::Retryable as i32
509        );
510        assert_eq!(terminal_wire.kind, ProtoActivityErrorKind::Terminal as i32);
511    }
512
513    #[tokio::test]
514    async fn typed_activity_round_trips_through_registry() -> Result<(), WorkerError> {
515        let registry =
516            ActivityRegistry::new().register_activity("double", |input: TestInput, context| {
517                Box::pin(async move {
518                    assert_eq!(context.attempt(), 1);
519                    Ok(TestOutput {
520                        doubled: input.value * 2,
521                    })
522                })
523            })?;
524        // One dispatch, one identity: the wire task and the handler context name
525        // the SAME workflow, run and activity. Independently minted ids would
526        // make the fixture blind to a stamper that dropped or crossed them.
527        let workflow_id = WorkflowId::new_v4();
528        let run_id = aion_core::RunId::new_v4();
529        let activity_id = ActivityId::from_sequence_position(99);
530        let task = proto_task(
531            "double",
532            &TestInput { value: 21 },
533            &workflow_id,
534            &run_id,
535            &activity_id,
536        )?;
537        let (context, cancellation) =
538            crate::ActivityContext::for_workflow(workflow_id, run_id, activity_id, 1, None);
539        drop(cancellation);
540
541        let outcome = registry.dispatch(task.try_into()?, context).await?;
542
543        let DispatchOutcome::Completed { output } = outcome else {
544            return Err(WorkerError::decode(UnexpectedFailure));
545        };
546        assert_eq!(output.content_type(), &ContentType::Json);
547        let decoded: TestOutput = decode_payload(&output)?;
548        assert_eq!(decoded, TestOutput { doubled: 42 });
549        Ok(())
550    }
551
552    #[test]
553    fn duplicate_activity_registration_is_rejected() -> Result<(), WorkerError> {
554        let registry =
555            ActivityRegistry::new().register_activity("double", |input: TestInput, context| {
556                Box::pin(async move {
557                    let _ = context;
558                    Ok(TestOutput {
559                        doubled: input.value * 2,
560                    })
561                })
562            })?;
563
564        let error = registry
565            .register_activity("double", |input: TestInput, context| {
566                Box::pin(async move {
567                    let _ = context;
568                    Ok(TestOutput {
569                        doubled: input.value,
570                    })
571                })
572            })
573            .err()
574            .ok_or_else(|| WorkerError::decode(UnexpectedFailure))?;
575
576        assert!(
577            error
578                .to_string()
579                .contains("already has a registered handler")
580        );
581        Ok(())
582    }
583
584    fn proto_task(
585        activity_type: &str,
586        input: &TestInput,
587        workflow_id: &WorkflowId,
588        run_id: &aion_core::RunId,
589        activity_id: &ActivityId,
590    ) -> Result<ProtoActivityTask, WorkerError> {
591        Ok(ProtoActivityTask {
592            workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
593            activity_id: Some(ProtoActivityId::from(activity_id.clone())),
594            run_id: Some(aion_proto::ProtoRunId::from(run_id.clone())),
595            activity_type: activity_type.to_owned(),
596            input: Some(ProtoPayload::from(encode_payload(&input)?)),
597            attempt: 1,
598            completion_token: String::from("generation-1"),
599            idempotency_key: String::from("effect-key"),
600            labels: std::collections::HashMap::new(),
601        })
602    }
603
604    #[derive(Debug, thiserror::Error)]
605    #[error("expected completed activity outcome")]
606    struct UnexpectedFailure;
607}
aion_worker/activity.rs

aion_worker/
activity.rs
