durable_macros/

lib.rs

use proc_macro::TokenStream;
use quote::quote;
use syn::{ItemFn, parse_macro_input};

/// Marks a function as a durable workflow with automatic crash recovery.
///
/// Supports two signatures:
///
/// **No input** — `fn(Ctx)`:
/// ```ignore
/// #[durable::workflow]
/// async fn daily_etl(ctx: Ctx) -> Result<EtlReport, DurableError> {
///     let data = ctx.step("extract", || async { fetch().await }).await?;
///     ctx.complete(&data).await?;
///     Ok(data)
/// }
///
/// // Generated: durable_daily_etl(&db, "my-etl").await?
/// ```
///
/// **Typed input** — `fn(Ctx, T)` where `T: Serialize + DeserializeOwned`:
/// ```ignore
/// #[derive(Serialize, Deserialize)]
/// struct IngestInput { crawl_id: String, shard_count: u32 }
///
/// #[durable::workflow]
/// async fn ingest(ctx: Ctx, input: IngestInput) -> Result<(), DurableError> {
///     for shard in 0..input.shard_count {
///         ctx.step(&format!("shard-{shard}"), || async {
///             process_shard(&input.crawl_id, shard).await
///         }).await?;
///     }
///     ctx.complete(&"done").await?;
///     Ok(())
/// }
///
/// // Generated: durable_ingest(&db, "my-ingest", IngestInput { crawl_id: "CC-2026".into(), shard_count: 72 }).await?
/// // On crash recovery: input is automatically deserialized and passed to ingest()
/// ```
#[proc_macro_attribute]
pub fn workflow(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input_fn = parse_macro_input!(item as ItemFn);
    let fn_name = &input_fn.sig.ident;
    let fn_name_str = fn_name.to_string();
    let vis = &input_fn.vis;
    let sig = &input_fn.sig;
    let body = &input_fn.block;
    let attrs = &input_fn.attrs;

    let is_ctx_type = |ty: &syn::Type| -> bool {
        matches!(ty, syn::Type::Path(tp) if tp.path.segments.last().is_some_and(|s| s.ident == "Ctx"))
    };

    // Detect signature shape
    let first_is_ctx = input_fn.sig.inputs.len() >= 1
        && matches!(&input_fn.sig.inputs[0], syn::FnArg::Typed(pat) if is_ctx_type(&pat.ty));

    let registration = if first_is_ctx && input_fn.sig.inputs.len() == 1 {
        // fn(Ctx) — no input, simple case
        let start_fn_name = syn::Ident::new(
            &format!("durable_{}", fn_name),
            fn_name.span(),
        );
        quote! {
            ::durable::inventory::submit! {
                ::durable::WorkflowRegistration {
                    name: #fn_name_str,
                    resume_fn: |ctx| ::std::boxed::Box::pin(async move {
                        let _ = #fn_name(ctx).await?;
                        Ok(())
                    }),
                }
            }

            /// Auto-generated by `#[durable::workflow]`. Starts (or idempotently
            /// attaches to) a root workflow, recording the handler name for crash
            /// recovery.
            #vis async fn #start_fn_name(
                db: &::sea_orm::DatabaseConnection,
                name: &str,
            ) -> ::std::result::Result<::durable::Ctx, ::durable::DurableError> {
                ::durable::Ctx::start_with_handler(db, name, ::std::option::Option::None, ::std::option::Option::Some(#fn_name_str)).await
            }
        }
    } else if first_is_ctx && input_fn.sig.inputs.len() == 2 {
        // fn(Ctx, T) — typed input
        let input_type = match &input_fn.sig.inputs[1] {
            syn::FnArg::Typed(pat) => &pat.ty,
            _ => panic!("#[durable::workflow] second parameter must be a typed argument"),
        };
        let start_fn_name = syn::Ident::new(
            &format!("durable_{}", fn_name),
            fn_name.span(),
        );
        quote! {
            ::durable::inventory::submit! {
                ::durable::WorkflowRegistration {
                    name: #fn_name_str,
                    resume_fn: |ctx| ::std::boxed::Box::pin(async move {
                        let input: #input_type = ctx.input().await?;
                        let _ = #fn_name(ctx, input).await?;
                        Ok(())
                    }),
                }
            }

            /// Auto-generated by `#[durable::workflow]`. Starts (or idempotently
            /// attaches to) a root workflow with typed input, recording the handler
            /// name for crash recovery.
            #vis async fn #start_fn_name(
                db: &::sea_orm::DatabaseConnection,
                name: &str,
                input: #input_type,
            ) -> ::std::result::Result<::durable::Ctx, ::durable::DurableError> {
                let input_json = ::serde_json::to_value(&input)
                    .map_err(|e| ::durable::DurableError::custom(format!("failed to serialize workflow input: {e}")))?;
                ::durable::Ctx::start_with_handler(db, name, ::std::option::Option::Some(input_json), ::std::option::Option::Some(#fn_name_str)).await
            }
        }
    } else {
        quote! {}
    };

    let expanded = quote! {
        #(#attrs)*
        #vis #sig {
            let _workflow_name = #fn_name_str;
            tracing::info!(workflow = _workflow_name, "workflow started");
            let _result: Result<_, _> = async { #body }.await;
            match &_result {
                Ok(_) => tracing::info!(workflow = _workflow_name, "workflow completed"),
                Err(e) => tracing::error!(workflow = _workflow_name, error = %e, "workflow failed"),
            }
            _result
        }

        #registration
    };

    TokenStream::from(expanded)
}

/// Marks a function as a durable step.
///
/// The first parameter must be `ctx: &Ctx`. The macro wraps the function
/// to check for saved output, execute if needed, and save the result.
///
/// ```ignore
/// #[durable::step]
/// async fn fetch_data(ctx: &Ctx, url: &str) -> Result<Data> {
///     // ...
/// }
/// ```
#[proc_macro_attribute]
pub fn step(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input_fn = parse_macro_input!(item as ItemFn);
    let fn_name = &input_fn.sig.ident;
    let fn_name_str = fn_name.to_string();
    let vis = &input_fn.vis;
    let sig = &input_fn.sig;
    let body = &input_fn.block;
    let attrs = &input_fn.attrs;

    let expanded = quote! {
        #(#attrs)*
        #vis #sig {
            let _step_name = #fn_name_str;
            tracing::debug!(step = _step_name, "step executing");
            let _result: Result<_, _> = async { #body }.await;
            match &_result {
                Ok(_) => tracing::debug!(step = _step_name, "step completed"),
                Err(e) => tracing::warn!(step = _step_name, error = %e, "step failed"),
            }
            _result
        }
    };

    TokenStream::from(expanded)
}