API improvements

- Added public `Error`, `Result`, `InitFuture`, and `ShutdownFuture`
  aliases.
- Added `TypeMapState` as the default state instead of `Arc<TypeMap>`.
- Added `App::route`, `App::mount`, and `App::store`.
- Changed `App::init()` to return `InitializedApp<S>`.
- Added `InitializedApp::router()`, `state()`, `into_parts()`, and
  `shutdown()`.
- Made shutdown hooks fallible: `Result<()>`.
- Made `AdHocPlugin::on_shutdown` accept capturing closures.
- Added `Default` for `App` and `AdHocPlugin`.

Improved AppState macro:
- Uses `axum_app_wrapper::Error`.
- Supports generic state structs.
- Fixed the stale `Arc<AppState>` doc snippet.

README.md

@@ -17,12 +17,32 @@ This crate is intentionally thin. It does not replace axum's router, extractors,
 1. `on_init` in registration order, passing a `TypeMap` to build state.
 2. `S::try_from(TypeMap)` to build the final typed app state.
 3. `on_setup` in registration order, passing `Router<S>` and `&S`.
-4. `on_shutdown` consecutively in reverse registration order: the last registered plugin shuts down first, and each hook finishes before the next one starts.
+4. `InitializedApp::shutdown()` runs `on_shutdown` consecutively in reverse registration order: the last registered plugin shuts down first, and each hook finishes before the next one starts.
 
 
 ## Example
 See the `examples` folder for full examples.
 
+`App::init()` returns an initialized app handle. Use `router()` to get a ready-to-serve `Router<()>`, `state()` to inspect the finalized state, and `shutdown()` from your graceful shutdown path:
+
+```rust
+let app = App::<AppState>::new()
+    .store(config)
+    .route("/health", axum::routing::get(health))
+    .register(metrics_plugin)
+    .init()
+    .await?;
+
+let router = app.router();
+let service_name = &app.state().config.service_name;
+
+axum::serve(listener, router)
+    .with_graceful_shutdown(async move {
+        tokio::signal::ctrl_c().await.expect("failed to listen for ctrl-c");
+        app.shutdown().await.expect("failed to shut down");
+    })
+    .await?;
+```
 
 For `on_setup` closures that access typed state, construct the plugin as
 `AdHocPlugin::<AppState>::new()`. That gives Rust enough context to infer the `state` parameter:
@@ -41,15 +61,15 @@ Implement `AppPlugin<S>` directly when setup should be reusable across apps:
 
 ```rust
 use axum::Router;
-use axum_app_wrapper::{AppPlugin, TypeMap};
+use axum_app_wrapper::{AppPlugin, InitFuture, Result, TypeMap};
-use futures::{future::BoxFuture, FutureExt};
+use futures::FutureExt;
 
 struct ConfigPlugin {
     config: Config,
 }
 
 impl AppPlugin<AppState> for ConfigPlugin {
-    fn on_init(&mut self, mut state: TypeMap) -> BoxFuture<'static, anyhow::Result<TypeMap>> {
+    fn on_init(&mut self, mut state: TypeMap) -> InitFuture {
         let config = self.config.clone();
         async move {
             state.insert(config);
@@ -62,7 +82,7 @@ impl AppPlugin<AppState> for ConfigPlugin {
         &mut self,
         router: Router<AppState>,
         _state: &AppState,
-    ) -> anyhow::Result<Router<AppState>> {
+    ) -> Result<Router<AppState>> {
         Ok(router.route("/health", axum::routing::get(health)))
     }
 }

crates/axum-app-wrapper-macros/src/lib.rs

@@ -22,10 +22,10 @@ use syn::{Data, DeriveInput, Fields, parse_macro_input};
 /// }
 /// // To wrap the whole state in `Arc`, implement `TryFrom<TypeMap>` for `Arc<AppState>`:
 /// impl TryFrom<TypeMap> for Arc<AppState> {
-///     type Error = anyhow::Error;
+///     type Error = axum_app_wrapper::Error;
 ///
 ///     fn try_from(map: TypeMap) -> Result<Self, Self::Error> {
-///         Ok(Self(Arc::new(AppState::try_from(map)?)))
+///         Ok(Arc::new(AppState::try_from(map)?))
 ///     }
 /// }
 /// ```
@@ -57,10 +57,11 @@ pub fn derive_app_state(input: TokenStream) -> TokenStream {
     };
 
     let field_names = fields.iter().map(|f| &f.ident);
+    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
 
     quote! {
-        impl ::std::convert::TryFrom<::axum_app_wrapper::TypeMap> for #name {
+        impl #impl_generics ::std::convert::TryFrom<::axum_app_wrapper::TypeMap> for #name #ty_generics #where_clause {
-            type Error = ::anyhow::Error;
+            type Error = ::axum_app_wrapper::Error;
 
             fn try_from(mut map: ::axum_app_wrapper::TypeMap) -> ::std::result::Result<Self, Self::Error> {
                 Ok(#name {

examples/app.rs

@@ -67,27 +67,30 @@ async fn main() -> anyhow::Result<()> {
             let metrics = Arc::clone(&state.metrics);
             async move {
                 metrics.flush().await;
+                Ok(())
             }
         });
 
-    // Register your plugins in the desired order
+    // Register your plugins in the desired order, and initialize the app
     let app = App::<AppState>::new()
         .register(config_plugin)
-        .register(metrics_plugin);
+        .register(metrics_plugin)
-
+        .init()
-    let (router, state, on_shutdown) = app.init().await?;
+        .await?;
-    let router = router.with_state(state);
+    tracing::info!(service = %app.state().config.service_name, "starting server");
 
     let addr: SocketAddr = "127.0.0.1:3000".parse()?;
     let listener = tokio::net::TcpListener::bind(addr).await?;
 
     // Start the axum server with graceful shutdown
-    axum::serve(listener, router)
+    axum::serve(listener, app.router())
         .with_graceful_shutdown(async {
             tokio::signal::ctrl_c()
                 .await
                 .expect("failed to listen for ctrl-c");
-            on_shutdown.await; // Run the on_shutdown future for graceful shutdown
+            app.shutdown()
+                .await
+                .expect("failed to run graceful shutdown");
         })
         .await?;
 

src/lib.rs

@@ -1,4 +1,3 @@
-#![allow(clippy::type_complexity)]
 //! A small plugin layer for [`axum`](https://docs.rs/axum/latest/axum/) applications.
 //!
 //! `axum-app-wrapper` lets plugins contribute startup state, router setup, and shutdown work
@@ -9,7 +8,7 @@
 //! 1. [`AppPlugin::on_init`] runs in registration order and builds a [`TypeMap`].
 //! 2. The `TypeMap` is converted into the app state `S` with `TryFrom<TypeMap>`.
 //! 3. [`AppPlugin::on_setup`] runs in registration order and receives `Router<S>` plus `&S`.
-//! 4. The returned shutdown future runs [`AppPlugin::on_shutdown`] consecutively in reverse
+//! 4. [`InitializedApp::shutdown`] runs [`AppPlugin::on_shutdown`] consecutively in reverse
 //!    registration order.
 //!
 //! For one-off plugins, use [`AdHocPlugin`]. If the setup closure needs access to typed state,
@@ -27,30 +26,60 @@
 //!     });
 //! ```
 
+extern crate self as axum_app_wrapper;
+
 use std::{fmt::Display, sync::Arc};
 
-use anyhow::anyhow;
+use axum::{Router, routing::MethodRouter};
-use axum::Router;
 use futures::{FutureExt, future::BoxFuture};
 
-// State extraction utilities
-
 pub use axum_app_wrapper_macros::AppState;
 pub use type_map::concurrent::TypeMap;
 
+pub use anyhow;
+
+/// Error type used by this crate.
+pub type Error = anyhow::Error;
+
+/// Result type used by this crate.
+pub type Result<T> = std::result::Result<T, Error>;
+
+/// Startup initialization future returned by [`AppPlugin::on_init`].
+pub type InitFuture = BoxFuture<'static, Result<TypeMap>>;
+
+/// Shutdown future returned by [`AppPlugin::on_shutdown`].
+pub type ShutdownFuture = BoxFuture<'static, Result<()>>;
+
+/// Default app state that keeps the startup [`TypeMap`] available behind an [`Arc`].
+#[derive(Clone)]
+pub struct TypeMapState(Arc<TypeMap>);
+
+impl TypeMapState {
+    /// Borrow the underlying [`TypeMap`].
+    pub fn type_map(&self) -> &TypeMap {
+        &self.0
+    }
+}
+
+impl From<TypeMap> for TypeMapState {
+    fn from(map: TypeMap) -> Self {
+        Self(Arc::new(map))
+    }
+}
+
 /// Extracts a value of type `T` from a [`TypeMap`], returning an Anyhow error if the type is not present.
 ///
 /// This is used internally by the [`AppState`] macro but is also available for manual
 /// [`TryFrom<TypeMap>`] implementations.
-pub fn extract_type_field<T: Send + Sync + 'static>(map: &mut TypeMap) -> anyhow::Result<T> {
+pub fn extract_type_field<T: Send + Sync + 'static>(map: &mut TypeMap) -> Result<T> {
     map.remove::<T>()
-        .ok_or_else(|| anyhow!("Missing type in TypeMap: {}", std::any::type_name::<T>()))
+        .ok_or_else(|| ::anyhow::anyhow!("Missing type in TypeMap: {}", std::any::type_name::<T>()))
 }
 
 // Plugin system
 
 /// An axum app builder with plugin-managed state, router setup, and shutdown hooks.
-pub struct App<S = Arc<TypeMap>> {
+pub struct App<S = TypeMapState> {
     base_router: Router<S>,
     plugins: Vec<Box<dyn AppPlugin<S>>>,
     state: TypeMap,
@@ -69,6 +98,27 @@ where
         }
     }
 
+    /// Add a route to the base router before plugins run setup.
+    pub fn route(mut self, path: &str, method_router: MethodRouter<S>) -> Self {
+        self.base_router = self.base_router.route(path, method_router);
+        self
+    }
+
+    /// Mount routes at the given path before plugins run setup.
+    pub fn mount(mut self, path: &str, router: Router<S>) -> Self {
+        self.base_router = match path {
+            "" | "/" => self.base_router.merge(router),
+            _ => self.base_router.nest(path, router),
+        };
+        self
+    }
+
+    /// Store a startup value for the final state conversion step.
+    pub fn store<T: Send + Sync + 'static>(mut self, state: T) -> Self {
+        self.state.insert(state);
+        self
+    }
+
     /// Register a plugin.
     ///
     /// `on_init` and `on_setup` run in registration order. Shutdown runs in reverse registration
@@ -78,26 +128,8 @@ where
         self
     }
 
-    // /// Mount the routes at the given path. All layers / middleware from plugins
+    /// Build the router, finalized state, and graceful shutdown handle.
-    // /// will be applied to these routes.
+    pub async fn init(mut self) -> Result<InitializedApp<S>>
-    // pub fn mount(mut self, path: &str, router: Router<S>) -> Self {
-    //     self.base_router = match path {
-    //         "" | "/" => self.base_router.merge(router),
-    //         _ => self.base_router.nest(path, router),
-    //     };
-    //     self
-    // }
-
-    // /// Store type T in state
-    // pub fn store<T: Send + Sync + 'static>(mut self, state: T) -> Self {
-    //     self.state.insert(state);
-    //     self
-    // }
-
-    /// Build the router, finalized state, and graceful shutdown future.
-    ///
-    /// The returned shutdown future must be awaited by the caller when the server is shutting down.
-    pub async fn init(mut self) -> anyhow::Result<(Router<S>, S, impl Future + Send)>
     where
         S::Error: Display,
     {
@@ -106,7 +138,8 @@ where
             state = plugin.on_init(state).await?;
         }
 
-        let state = S::try_from(state).map_err(|err| anyhow!("Error creating state: {err}"))?;
+        let state =
+            S::try_from(state).map_err(|err| ::anyhow::anyhow!("Error creating state: {err}"))?;
 
         let mut router = self.base_router;
         for plugin in self.plugins.iter_mut() {
@@ -121,35 +154,82 @@ where
             .collect();
         let on_shutdown = async move {
             for shutdown_fn in shutdown_fns {
-                shutdown_fn.await;
+                shutdown_fn.await?;
             }
-        };
+            Ok(())
+        }
+        .boxed();
 
-        Ok((router, state, on_shutdown))
+        Ok(InitializedApp {
+            router,
+            state,
+            on_shutdown,
+        })
+    }
+}
+
+impl<S> Default for App<S>
+where
+    S: TryFrom<TypeMap> + Clone + Send + Sync + 'static,
+{
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// A fully initialized axum app.
+pub struct InitializedApp<S> {
+    router: Router<S>,
+    state: S,
+    on_shutdown: ShutdownFuture,
+}
+
+impl<S> InitializedApp<S>
+where
+    S: Clone + Send + Sync + 'static,
+{
+    /// Build a ready-to-serve router by attaching the finalized state.
+    pub fn router(&self) -> Router<()> {
+        self.router.clone().with_state(self.state.clone())
+    }
+
+    /// Borrow the finalized app state.
+    pub fn state(&self) -> &S {
+        &self.state
+    }
+
+    /// Consume the initialized app and return its raw parts.
+    pub fn into_parts(self) -> (Router<S>, S, ShutdownFuture) {
+        (self.router, self.state, self.on_shutdown)
+    }
+
+    /// Run graceful shutdown work for all plugins.
+    pub async fn shutdown(self) -> Result<()> {
+        self.on_shutdown.await
     }
 }
 
 /// A plugin that can participate in app initialization, router setup, and shutdown.
 #[allow(unused_variables, reason = "trait functions with default no-op")]
-pub trait AppPlugin<S = Arc<TypeMap>> {
+pub trait AppPlugin<S = TypeMapState> {
     /// Run during startup before typed state exists.
     ///
     /// Use this hook to insert or transform values in the shared [`TypeMap`].
-    fn on_init(&mut self, app_state: TypeMap) -> BoxFuture<'static, anyhow::Result<TypeMap>> {
+    fn on_init(&mut self, app_state: TypeMap) -> InitFuture {
         async { Ok(app_state) }.boxed()
     }
 
     /// Run after typed state has been created.
     ///
     /// Use this hook to add routes, services, middleware, or router state.
-    fn on_setup(&mut self, router: Router<S>, state: &S) -> anyhow::Result<Router<S>> {
+    fn on_setup(&mut self, router: Router<S>, state: &S) -> Result<Router<S>> {
         Ok(router)
     }
 
     /// Return shutdown work for this plugin.
     ///
     /// Shutdown hooks are awaited consecutively in reverse registration order.
-    fn on_shutdown(&mut self, state: &S) -> Option<BoxFuture<'static, ()>> {
+    fn on_shutdown(&mut self, state: &S) -> Option<ShutdownFuture> {
         None
     }
 }
@@ -158,12 +238,16 @@ pub trait AppPlugin<S = Arc<TypeMap>> {
 ///
 /// `on_init` and `on_setup` accept capturing closures. If `on_setup` uses typed state, prefer
 /// `AdHocPlugin::<State>::new()` so the closure parameter type can be inferred.
-pub struct AdHocPlugin<S = Arc<TypeMap>> {
+pub struct AdHocPlugin<S = TypeMapState> {
-    on_init: Option<Box<dyn FnOnce(TypeMap) -> BoxFuture<'static, anyhow::Result<TypeMap>> + Send>>,
+    on_init: Option<InitFn>,
-    on_setup: Option<Box<dyn FnOnce(Router<S>, &S) -> anyhow::Result<Router<S>> + Send>>,
+    on_setup: Option<SetupFn<S>>,
-    on_shutdown: Option<Box<dyn FnOnce(&S) -> BoxFuture<'static, ()>>>,
+    on_shutdown: Option<ShutdownFn<S>>,
 }
 
+type InitFn = Box<dyn FnOnce(TypeMap) -> InitFuture + Send>;
+type SetupFn<S> = Box<dyn FnOnce(Router<S>, &S) -> Result<Router<S>> + Send>;
+type ShutdownFn<S> = Box<dyn FnOnce(&S) -> ShutdownFuture + Send>;
+
 impl<S: 'static> AdHocPlugin<S> {
     /// Create an empty ad-hoc plugin.
     pub fn new() -> Self {
@@ -178,7 +262,7 @@ impl<S: 'static> AdHocPlugin<S> {
     pub fn on_init<F, T>(mut self, on_init: F) -> Self
     where
         F: FnOnce(TypeMap) -> T + Send + 'static,
-        T: Future<Output = anyhow::Result<TypeMap>> + Send + 'static,
+        T: Future<Output = Result<TypeMap>> + Send + 'static,
     {
         self.on_init = Some(Box::new(move |s| Box::pin(on_init(s))));
         self
@@ -187,38 +271,45 @@ impl<S: 'static> AdHocPlugin<S> {
     /// Set router setup for this plugin.
     pub fn on_setup<F>(mut self, on_setup: F) -> Self
     where
-        F: FnOnce(Router<S>, &S) -> anyhow::Result<Router<S>> + Send + 'static,
+        F: FnOnce(Router<S>, &S) -> Result<Router<S>> + Send + 'static,
     {
         self.on_setup = Some(Box::new(on_setup));
         self
     }
 
     /// Set shutdown work for this plugin.
-    pub fn on_shutdown<T>(mut self, on_shutdown: fn(state: &S) -> T) -> Self
+    pub fn on_shutdown<F, T>(mut self, on_shutdown: F) -> Self
     where
-        T: Future<Output = ()> + Send + 'static,
+        F: FnOnce(&S) -> T + Send + 'static,
+        T: Future<Output = Result<()>> + Send + 'static,
     {
         self.on_shutdown = Some(Box::new(move |s| Box::pin(on_shutdown(s))));
         self
     }
 }
 
+impl<S: 'static> Default for AdHocPlugin<S> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl<S> AppPlugin<S> for AdHocPlugin<S> {
-    fn on_init(&mut self, app_state: TypeMap) -> BoxFuture<'static, anyhow::Result<TypeMap>> {
+    fn on_init(&mut self, app_state: TypeMap) -> InitFuture {
         match self.on_init.take() {
             Some(init_fn) => async move { init_fn(app_state).await }.boxed(),
             None => async { Ok(app_state) }.boxed(),
         }
     }
 
-    fn on_setup(&mut self, router: Router<S>, state: &S) -> anyhow::Result<Router<S>> {
+    fn on_setup(&mut self, router: Router<S>, state: &S) -> Result<Router<S>> {
         match self.on_setup.take() {
             Some(setup_fn) => setup_fn(router, state),
             None => Ok(router),
         }
     }
 
-    fn on_shutdown(&mut self, state: &S) -> Option<BoxFuture<'static, ()>> {
+    fn on_shutdown(&mut self, state: &S) -> Option<ShutdownFuture> {
         match self.on_shutdown.take() {
             Some(shutdown_fn) => Some(shutdown_fn(state)),
             None => None,
@@ -228,8 +319,9 @@ impl<S> AppPlugin<S> for AdHocPlugin<S> {
 
 #[cfg(test)]
 mod tests {
+    use super::*;
+
     use std::{
-        convert::Infallible,
         sync::{
             Mutex,
             atomic::{AtomicUsize, Ordering},
@@ -237,23 +329,11 @@ mod tests {
         task::Poll,
     };
 
-    use super::*;
+    #[derive(Clone, self::AppState)]
-
-    #[derive(Clone)]
     struct TestState {
         value: Arc<String>,
     }
 
-    impl TryFrom<TypeMap> for TestState {
-        type Error = Infallible;
-
-        fn try_from(mut map: TypeMap) -> Result<Self, Self::Error> {
-            Ok(Self {
-                value: map.remove::<Arc<String>>().expect("missing state value"),
-            })
-        }
-    }
-
     #[test]
     fn adhoc_plugin_with_basic_state() {
         let init_value = Arc::new(String::from("ready"));
@@ -270,10 +350,75 @@ mod tests {
             });
         let app = App::<TestState>::new().register(plugin);
 
-        let (_router, state, _shutdown) =
+        let app = futures::executor::block_on(app.init()).expect("app should initialize");
-            futures::executor::block_on(app.init()).expect("app should initialize");
 
-        assert_eq!(state.value.as_str(), "ready");
+        assert_eq!(app.state().value.as_str(), "ready");
+    }
+
+    #[test]
+    fn store_and_router_handle_expose_initialized_state() {
+        let app = App::<TestState>::new().store(Arc::new(String::from("stored")));
+
+        let app = futures::executor::block_on(app.init()).expect("app should initialize");
+        let _router = app.router();
+
+        assert_eq!(app.state().value.as_str(), "stored");
+    }
+
+    #[test]
+    fn default_state_keeps_type_map_available() {
+        let app: App = App::new().store(Arc::new(String::from("typemap")));
+
+        let app = futures::executor::block_on(app.init()).expect("app should initialize");
+
+        let value = app
+            .state()
+            .type_map()
+            .get::<Arc<String>>()
+            .expect("stored value should remain in typemap");
+        assert_eq!(value.as_str(), "typemap");
+    }
+
+    #[derive(Clone, AppState)]
+    struct GenericState<T: Send + Sync + 'static> {
+        value: Arc<T>,
+    }
+
+    #[test]
+    fn app_state_derive_supports_generics() {
+        let app = App::<GenericState<String>>::new().store(Arc::new(String::from("generic")));
+
+        let app = futures::executor::block_on(app.init()).expect("app should initialize");
+
+        assert_eq!(app.state().value.as_str(), "generic");
+    }
+
+    #[test]
+    fn adhoc_shutdown_accepts_capturing_fallible_closure() {
+        let events = Arc::new(Mutex::new(Vec::new()));
+        let shutdown_events = Arc::clone(&events);
+
+        let app = App::<TestState>::new()
+            .store(Arc::new(String::from("ready")))
+            .register(AdHocPlugin::<TestState>::new().on_shutdown(move |state| {
+                let events = Arc::clone(&shutdown_events);
+                let value = Arc::clone(&state.value);
+                async move {
+                    events
+                        .lock()
+                        .expect("events lock poisoned")
+                        .push(value.to_string());
+                    Ok(())
+                }
+            }));
+
+        let app = futures::executor::block_on(app.init()).expect("app should initialize");
+        futures::executor::block_on(app.shutdown()).expect("shutdown should succeed");
+
+        assert_eq!(
+            *events.lock().expect("events lock poisoned"),
+            [String::from("ready")]
+        );
     }
 
     struct ShutdownOrderPlugin {
@@ -283,7 +428,7 @@ mod tests {
     }
 
     impl AppPlugin<TestState> for ShutdownOrderPlugin {
-        fn on_shutdown(&mut self, _state: &TestState) -> Option<BoxFuture<'static, ()>> {
+        fn on_shutdown(&mut self, _state: &TestState) -> Option<ShutdownFuture> {
             let name = self.name;
             let events = Arc::clone(&self.events);
             let active_shutdowns = Arc::clone(&self.active_shutdowns);
@@ -307,7 +452,7 @@ mod tests {
                     .expect("events lock poisoned")
                     .push(format!("{name}:finish"));
                 active_shutdowns.fetch_sub(1, Ordering::SeqCst);
-                Poll::Ready(())
+                Poll::Ready(Ok(()))
             })))
         }
     }
@@ -333,9 +478,8 @@ mod tests {
                 active_shutdowns,
             });
 
-        let (_router, _state, on_shutdown) =
+        let app = futures::executor::block_on(app.init()).expect("app should initialize");
-            futures::executor::block_on(app.init()).expect("app should initialize");
+        futures::executor::block_on(app.shutdown()).expect("shutdown should succeed");
-        futures::executor::block_on(on_shutdown);
 
         assert_eq!(
             *events.lock().expect("events lock poisoned"),