Expose the web server as an embeddable Axum router

2026-09-15 22:32:04 +02:00 · 2026-09-14 22:21:43 -07:00
parent ff5db718ce
commit 8c8be31aa2
6 changed files with 381 additions and 229 deletions
@@ -7,7 +7,7 @@ description = "The server behind Yaak in the browser: executes sends, and can se

 # The send engine (yaak-http) and the model types it speaks (yaak-models, for
 # HttpRequest / Cookie / HttpResponseEventData). Deliberately NOT yaak (the
-# render + storage orchestration), yaak-plugins, or the RPC router: this binary
+# render + storage orchestration), yaak-plugins, or the RPC router: this crate
 # opens no database, runs no plugins, and renders nothing. yaak-models comes
 # along only because yaak-http's types are its types; nothing here calls into
 # its query layer.
@@ -34,3 +34,7 @@ url = "2"
 uuid = { version = "1", features = ["v4"] }
 yaak-http = { workspace = true }
 yaak-models = { workspace = true }
+
+[dev-dependencies]
+tempfile = "3"
+tower = { version = "0.5", features = ["util"] }
@@ -76,6 +76,50 @@ A dev build looks for the server at `http://127.0.0.1:9227` (the Vite server is
 different origin and serves no `/v1`); a production build sends to its own
 origin unless `VITE_YAAK_WEB_URL` was set when it was built.

+## Embedding in your own server
+
+The crate also exposes `yaak_web::router(Config) -> axum::Router`. It builds the
+same API and optional static-file server as the binary, with its state already
+attached. Your server owns logging, the listener, shutdown, and any additional
+middleware. Deployment-specific integrations can live entirely in that server.
+
+Use Axum 0.7, matching this crate. For example:
+
+```rust,no_run
+use axum::{extract::Request, middleware::{self, Next}, response::Response};
+use clap::Parser;
+use std::net::SocketAddr;
+use yaak_web::Config;
+
+async fn hosted_middleware(request: Request, next: Next) -> Response {
+    // Add your hosted server's request/response handling here.
+    next.run(request).await
+}
+
+#[tokio::main]
+async fn main() -> std::io::Result<()> {
+    let config = Config::parse();
+    let listener = tokio::net::TcpListener::bind(config.bind).await?;
+    let app = yaak_web::router(config)
+        .layer(middleware::from_fn(hosted_middleware));
+
+    axum::serve(
+        listener,
+        app.into_make_service_with_connect_info::<SocketAddr>(),
+    )
+    .with_graceful_shutdown(async { tokio::signal::ctrl_c().await.ok(); })
+    .await
+}
+```
+
+You can also construct `Config` directly; the library does not read arguments or
+environment variables itself. Its `bind` field is only a convenience for callers
+that create a listener. Keep `into_make_service_with_connect_info` when serving:
+the send endpoint requires the socket peer address for rate limiting.
+
+Use `.layer()` to wrap both the API and static-file fallback. `.route_layer()`
+does not wrap the fallback, which serves the HTML pages and assets.
+
 ## Configuration

 Every flag has a `YAAK_WEB_*` environment variable, so a container needs no
@@ -162,7 +206,7 @@ other side is one the users are entitled to.
 There is no authentication either way: an instance is anonymous, protected by
 the per-client rate limit and the destination policy. Anything more (a shared
 token, per-user quotas) is a later slice and would sit in front of `send_http`
-in `main.rs`. Put TLS in front of a public instance.
+in `lib.rs`, or in middleware on the exported router. Put TLS in front of a public instance.

 ## The wire

@@ -0,0 +1,234 @@
+//! yaak-web: the network half of Yaak in a browser.
+//!
+//! A tab can't see a response the way a desktop app can — CORS hides most
+//! headers, redirects are followed silently, there is no timeline. So the tab
+//! renders the request and hands it here; this process puts it on the network
+//! with the desktop's own engine and streams back everything that happened,
+//! for the tab to store. It keeps nothing: no database, no files, no session.
+//!
+//! Embed [`router`] in your own Axum server, or run the `yaak-web` binary
+//! configured by flags or `YAAK_WEB_*` environment variables.
+//! See README.md for running and deploying it, and `guard.rs` for what it
+//! refuses to talk to.
+
+mod config;
+mod guard;
+mod limits;
+mod send;
+mod wire;
+
+use axum::Router;
+use axum::body::Body;
+use axum::extract::{ConnectInfo, DefaultBodyLimit, Request, State};
+use axum::http::{HeaderMap, HeaderValue, Method, StatusCode, header};
+use axum::middleware::{self, Next};
+use axum::response::{IntoResponse, Json, Response};
+use axum::routing::{get, post};
+pub use config::Config;
+use guard::DestinationPolicy;
+use limits::RateLimiter;
+use log::{info, warn};
+use send::{Refusal, SendLimits};
+use serde_json::json;
+use std::net::{IpAddr, SocketAddr};
+use std::path::Path;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+use tokio::sync::Semaphore;
+use tower_http::compression::CompressionLayer;
+use tower_http::cors::{AllowOrigin, CorsLayer};
+use tower_http::services::{ServeDir, ServeFile};
+use wire::SendRequest;
+
+#[derive(Clone)]
+struct AppState {
+    config: Arc<Config>,
+    limits: Arc<SendLimits>,
+    rate_limiter: Arc<RateLimiter>,
+    in_flight: Arc<Semaphore>,
+}
+
+/// Build the API router and, when configured, the web client's static-file fallback.
+///
+/// The returned router owns its state and accepts normal Axum middleware via `.layer()`.
+/// Construction does not parse arguments, initialize logging, bind a socket, or install
+/// shutdown handlers. `config.bind` is used only by the standalone binary.
+///
+/// Serve with `into_make_service_with_connect_info::<SocketAddr>()` so the send endpoint
+/// can extract the peer address for rate limiting.
+pub fn router(config: Config) -> Router {
+    let policy = DestinationPolicy::new(config.allow_private_networks);
+    if config.allow_private_networks {
+        warn!(
+            "Sends to loopback, private and link-local addresses are ALLOWED. Only run this way \
+             on an instance strangers cannot reach"
+        );
+    }
+    let state = AppState {
+        limits: Arc::new(SendLimits {
+            policy,
+            max_response_bytes: config.max_response_bytes,
+            max_timeout: Duration::from_secs(config.max_timeout_secs),
+        }),
+        rate_limiter: Arc::new(RateLimiter::new(config.rate_limit_per_minute)),
+        in_flight: Arc::new(Semaphore::new(config.max_concurrent)),
+        config: Arc::new(config),
+    };
+
+    let cors = CorsLayer::new()
+        .allow_methods([Method::GET, Method::POST, Method::OPTIONS])
+        .allow_headers([header::CONTENT_TYPE])
+        .allow_origin(allowed_origins(&state.config.allowed_origins));
+
+    let api = Router::new()
+        .route("/v1/health", get(health))
+        // A WebSocket or gRPC relay would sit beside this as `/v1/ws/relay` and `/v1/grpc/relay`
+        // on the same router, behind the same policy, limits and auth. Not built; see README.
+        .route("/v1/http/send", post(send_http))
+        .layer(DefaultBodyLimit::max(state.config.max_request_bytes))
+        .layer(cors)
+        .with_state(state.clone());
+
+    match &state.config.serve {
+        Some(dir) => {
+            info!("Serving the web client from {}", dir.display());
+            api.merge(web_router(dir))
+        }
+        None => api,
+    }
+}
+
+/// The built web client, served on the same origin as the API.
+///
+/// This is what makes a single container zero-configuration: the tab's send URL is a path on
+/// the page's own origin, so there is no CORS, no second service and no URL to bake in. It is
+/// only a file server — a send behaves exactly as it does without this flag.
+///
+/// Merged as a fallback, so the `/v1` routes are matched first and a request that matches no
+/// file at all gets `index.html` (the app routes client-side; a deep link must survive a
+/// refresh).
+fn web_router(dir: &Path) -> Router {
+    let index = ServeFile::new(dir.join("index.html"));
+    Router::new()
+        // `fallback`, not `not_found_service`: the app's own routes are real pages, so
+        // index.html is served with the 200 the browser expects, not a 404 carrying HTML.
+        .fallback_service(ServeDir::new(dir).fallback(index))
+        .layer(middleware::from_fn(cache_control))
+        .layer(CompressionLayer::new())
+}
+
+/// Vite gives everything in `/assets` a content-hashed name, so those can be cached forever.
+/// Everything else — `index.html` above all, including the copy served for an unknown path —
+/// must be revalidated, or a browser keeps serving the deploy before last.
+async fn cache_control(req: Request, next: Next) -> Response {
+    let hashed_name = req.uri().path().starts_with("/assets/");
+    let mut res = next.run(req).await;
+    if !res.status().is_success() {
+        return res;
+    }
+    let is_html = res
+        .headers()
+        .get(header::CONTENT_TYPE)
+        .and_then(|v| v.to_str().ok())
+        .is_some_and(|v| v.starts_with("text/html"));
+    let value =
+        if hashed_name && !is_html { "public, max-age=31536000, immutable" } else { "no-cache" };
+    res.headers_mut().insert(header::CACHE_CONTROL, HeaderValue::from_static(value));
+    res
+}
+
+fn allowed_origins(origins: &[String]) -> AllowOrigin {
+    if origins.iter().any(|o| o.trim() == "*") {
+        return AllowOrigin::any();
+    }
+    let parsed: Vec<HeaderValue> =
+        origins.iter().filter_map(|o| HeaderValue::from_str(o.trim()).ok()).collect();
+    AllowOrigin::list(parsed)
+}
+
+async fn health(State(state): State<AppState>) -> impl IntoResponse {
+    Json(json!({
+        "ok": true,
+        "version": env!("CARGO_PKG_VERSION"),
+        "maxResponseBytes": state.config.max_response_bytes,
+        "maxTimeoutSecs": state.config.max_timeout_secs,
+    }))
+}
+
+fn error_response(status: StatusCode, message: impl Into<String>) -> Response {
+    let message = message.into();
+    (status, Json(json!({ "error": message }))).into_response()
+}
+
+/// The client's address for rate limiting: the socket peer, or the first `X-Forwarded-For`
+/// hop when the operator has said the header can be trusted.
+fn client_ip(config: &Config, headers: &HeaderMap, peer: SocketAddr) -> IpAddr {
+    if config.trust_forwarded_for
+        && let Some(forwarded) = headers.get("x-forwarded-for").and_then(|v| v.to_str().ok())
+        && let Some(first) = forwarded.split(',').next()
+        && let Ok(ip) = first.trim().parse::<IpAddr>()
+    {
+        return ip;
+    }
+    peer.ip()
+}
+
+async fn send_http(
+    State(state): State<AppState>,
+    ConnectInfo(peer): ConnectInfo<SocketAddr>,
+    headers: HeaderMap,
+    Json(body): Json<SendRequest>,
+) -> Response {
+    let ip = client_ip(&state.config, &headers, peer);
+    if let Err(wait) = state.rate_limiter.check(ip) {
+        warn!("Rate limited {ip}");
+        let mut res = error_response(
+            StatusCode::TOO_MANY_REQUESTS,
+            format!("Rate limit reached; try again in {}s", wait.as_secs().max(1)),
+        );
+        res.headers_mut().insert(header::RETRY_AFTER, HeaderValue::from(wait.as_secs().max(1)));
+        return res;
+    }
+
+    let Ok(permit) = state.in_flight.clone().try_acquire_owned() else {
+        warn!("At capacity; refusing {ip}");
+        return error_response(StatusCode::SERVICE_UNAVAILABLE, "This server is at capacity");
+    };
+
+    let prepared = match send::prepare(state.limits.clone(), body).await {
+        Ok(p) => p,
+        Err(Refusal::Unsupported(m)) => return error_response(StatusCode::BAD_REQUEST, m),
+        Err(Refusal::Invalid(m)) => return error_response(StatusCode::BAD_REQUEST, m),
+        Err(Refusal::Destination(m)) => {
+            warn!("Refused send from {ip}: {m}");
+            return error_response(StatusCode::FORBIDDEN, m);
+        }
+    };
+
+    let description = prepared.describe();
+    info!("{ip} -> {description}");
+    let started = Instant::now();
+
+    let (tx, rx) = tokio::sync::mpsc::channel(send::FRAME_CHANNEL_CAPACITY);
+    tokio::spawn(async move {
+        prepared.run(tx).await;
+        send::log_outcome(&description, started, "finished");
+        drop(permit);
+    });
+
+    let stream = tokio_stream_from(rx);
+    Response::builder()
+        .status(StatusCode::OK)
+        .header(header::CONTENT_TYPE, "application/x-ndjson")
+        .header(header::CACHE_CONTROL, "no-store")
+        // Some reverse proxies buffer streamed responses unless told not to
+        .header("x-accel-buffering", "no")
+        .body(Body::from_stream(stream))
+        .expect("valid response")
+}
+
+fn tokio_stream_from<T: Send + 'static>(
+    mut rx: tokio::sync::mpsc::Receiver<T>,
+) -> impl futures_util::Stream<Item = T> + Send + 'static {
+    futures_util::stream::poll_fn(move |cx| rx.poll_recv(cx))
+}
@@ -1,107 +1,23 @@
-//! yaak-web: the network half of Yaak in a browser.
-//!
-//! A tab can't see a response the way a desktop app can — CORS hides most
-//! headers, redirects are followed silently, there is no timeline. So the tab
-//! renders the request and hands it here; this process puts it on the network
-//! with the desktop's own engine and streams back everything that happened,
-//! for the tab to store. It keeps nothing: no database, no files, no session.
-//!
-//! One binary, configured by flags or `YAAK_WEB_*` environment variables.
-//! See README.md for running and deploying it, and `guard.rs` for what it
-//! refuses to talk to.
+//! Standalone server for Yaak in a browser.

-mod config;
-mod guard;
-mod limits;
-mod send;
-mod wire;
-
-use axum::Router;
-use axum::body::Body;
-use axum::extract::{ConnectInfo, DefaultBodyLimit, Request, State};
-use axum::http::{HeaderMap, HeaderValue, Method, StatusCode, header};
-use axum::middleware::{self, Next};
-use axum::response::{IntoResponse, Json, Response};
-use axum::routing::{get, post};
 use clap::Parser;
-use config::Config;
-use guard::DestinationPolicy;
-use limits::RateLimiter;
-use log::{info, warn};
-use send::{Refusal, SendLimits};
-use serde_json::json;
-use std::net::{IpAddr, SocketAddr};
-use std::path::Path;
-use std::sync::Arc;
-use std::time::{Duration, Instant};
-use tokio::sync::Semaphore;
-use tower_http::compression::CompressionLayer;
-use tower_http::cors::{AllowOrigin, CorsLayer};
-use tower_http::services::{ServeDir, ServeFile};
-use wire::SendRequest;
-
-#[derive(Clone)]
-struct AppState {
-    config: Arc<Config>,
-    limits: Arc<SendLimits>,
-    rate_limiter: Arc<RateLimiter>,
-    in_flight: Arc<Semaphore>,
-}
+use log::info;
+use std::net::SocketAddr;
+use yaak_web::{Config, router};

 #[tokio::main]
 async fn main() {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
    let config = Config::parse();
+    let bind = config.bind;
+    let rate_limit_per_minute = config.rate_limit_per_minute;
+    let app = router(config);

-    let policy = DestinationPolicy::new(config.allow_private_networks);
-    if config.allow_private_networks {
-        warn!(
-            "Sends to loopback, private and link-local addresses are ALLOWED. Only run this way \
-             on an instance strangers cannot reach"
-        );
-    }
-    let state = AppState {
-        limits: Arc::new(SendLimits {
-            policy,
-            max_response_bytes: config.max_response_bytes,
-            max_timeout: Duration::from_secs(config.max_timeout_secs),
-        }),
-        rate_limiter: Arc::new(RateLimiter::new(config.rate_limit_per_minute)),
-        in_flight: Arc::new(Semaphore::new(config.max_concurrent)),
-        config: Arc::new(config),
-    };
-
-    let cors = CorsLayer::new()
-        .allow_methods([Method::GET, Method::POST, Method::OPTIONS])
-        .allow_headers([header::CONTENT_TYPE])
-        .allow_origin(allowed_origins(&state.config.allowed_origins));
-
-    let api = Router::new()
-        .route("/v1/health", get(health))
-        // A WebSocket or gRPC relay would sit beside this as `/v1/ws/relay` and `/v1/grpc/relay`
-        // on the same router, behind the same policy, limits and auth. Not built; see README.
-        .route("/v1/http/send", post(send_http))
-        .layer(DefaultBodyLimit::max(state.config.max_request_bytes))
-        .layer(cors)
-        .with_state(state.clone());
-
-    let app = match &state.config.serve {
-        Some(dir) => {
-            info!("Serving the web client from {}", dir.display());
-            api.merge(web_router(dir))
-        }
-        None => api,
-    };
-
-    let bind = state.config.bind;
    let listener = tokio::net::TcpListener::bind(bind).await.unwrap_or_else(|e| {
        eprintln!("Failed to bind {bind}: {e}");
        std::process::exit(1);
    });
-    info!(
-        "yaak-web listening on http://{bind} (rate limit: {}/min)",
-        state.config.rate_limit_per_minute,
-    );
+    info!("yaak-web listening on http://{bind} (rate limit: {}/min)", rate_limit_per_minute,);

    axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>())
        .with_graceful_shutdown(async {
@@ -111,138 +27,3 @@ async fn main() {
        .await
        .expect("server error");
 }
-
-/// The built web client, served on the same origin as the API.
-///
-/// This is what makes a single container zero-configuration: the tab's send URL is a path on
-/// the page's own origin, so there is no CORS, no second service and no URL to bake in. It is
-/// only a file server — a send behaves exactly as it does without this flag.
-///
-/// Merged as a fallback, so the `/v1` routes are matched first and a request that matches no
-/// file at all gets `index.html` (the app routes client-side; a deep link must survive a
-/// refresh).
-fn web_router(dir: &Path) -> Router {
-    let index = ServeFile::new(dir.join("index.html"));
-    Router::new()
-        // `fallback`, not `not_found_service`: the app's own routes are real pages, so
-        // index.html is served with the 200 the browser expects, not a 404 carrying HTML.
-        .fallback_service(ServeDir::new(dir).fallback(index))
-        .layer(middleware::from_fn(cache_control))
-        .layer(CompressionLayer::new())
-}
-
-/// Vite gives everything in `/assets` a content-hashed name, so those can be cached forever.
-/// Everything else — `index.html` above all, including the copy served for an unknown path —
-/// must be revalidated, or a browser keeps serving the deploy before last.
-async fn cache_control(req: Request, next: Next) -> Response {
-    let hashed_name = req.uri().path().starts_with("/assets/");
-    let mut res = next.run(req).await;
-    if !res.status().is_success() {
-        return res;
-    }
-    let is_html = res
-        .headers()
-        .get(header::CONTENT_TYPE)
-        .and_then(|v| v.to_str().ok())
-        .is_some_and(|v| v.starts_with("text/html"));
-    let value =
-        if hashed_name && !is_html { "public, max-age=31536000, immutable" } else { "no-cache" };
-    res.headers_mut().insert(header::CACHE_CONTROL, HeaderValue::from_static(value));
-    res
-}
-
-fn allowed_origins(origins: &[String]) -> AllowOrigin {
-    if origins.iter().any(|o| o.trim() == "*") {
-        return AllowOrigin::any();
-    }
-    let parsed: Vec<HeaderValue> =
-        origins.iter().filter_map(|o| HeaderValue::from_str(o.trim()).ok()).collect();
-    AllowOrigin::list(parsed)
-}
-
-async fn health(State(state): State<AppState>) -> impl IntoResponse {
-    Json(json!({
-        "ok": true,
-        "version": env!("CARGO_PKG_VERSION"),
-        "maxResponseBytes": state.config.max_response_bytes,
-        "maxTimeoutSecs": state.config.max_timeout_secs,
-    }))
-}
-
-fn error_response(status: StatusCode, message: impl Into<String>) -> Response {
-    let message = message.into();
-    (status, Json(json!({ "error": message }))).into_response()
-}
-
-/// The client's address for rate limiting: the socket peer, or the first `X-Forwarded-For`
-/// hop when the operator has said the header can be trusted.
-fn client_ip(config: &Config, headers: &HeaderMap, peer: SocketAddr) -> IpAddr {
-    if config.trust_forwarded_for
-        && let Some(forwarded) = headers.get("x-forwarded-for").and_then(|v| v.to_str().ok())
-        && let Some(first) = forwarded.split(',').next()
-        && let Ok(ip) = first.trim().parse::<IpAddr>()
-    {
-        return ip;
-    }
-    peer.ip()
-}
-
-async fn send_http(
-    State(state): State<AppState>,
-    ConnectInfo(peer): ConnectInfo<SocketAddr>,
-    headers: HeaderMap,
-    Json(body): Json<SendRequest>,
-) -> Response {
-    let ip = client_ip(&state.config, &headers, peer);
-    if let Err(wait) = state.rate_limiter.check(ip) {
-        warn!("Rate limited {ip}");
-        let mut res = error_response(
-            StatusCode::TOO_MANY_REQUESTS,
-            format!("Rate limit reached; try again in {}s", wait.as_secs().max(1)),
-        );
-        res.headers_mut().insert(header::RETRY_AFTER, HeaderValue::from(wait.as_secs().max(1)));
-        return res;
-    }
-
-    let Ok(permit) = state.in_flight.clone().try_acquire_owned() else {
-        warn!("At capacity; refusing {ip}");
-        return error_response(StatusCode::SERVICE_UNAVAILABLE, "This server is at capacity");
-    };
-
-    let prepared = match send::prepare(state.limits.clone(), body).await {
-        Ok(p) => p,
-        Err(Refusal::Unsupported(m)) => return error_response(StatusCode::BAD_REQUEST, m),
-        Err(Refusal::Invalid(m)) => return error_response(StatusCode::BAD_REQUEST, m),
-        Err(Refusal::Destination(m)) => {
-            warn!("Refused send from {ip}: {m}");
-            return error_response(StatusCode::FORBIDDEN, m);
-        }
-    };
-
-    let description = prepared.describe();
-    info!("{ip} -> {description}");
-    let started = Instant::now();
-
-    let (tx, rx) = tokio::sync::mpsc::channel(send::FRAME_CHANNEL_CAPACITY);
-    tokio::spawn(async move {
-        prepared.run(tx).await;
-        send::log_outcome(&description, started, "finished");
-        drop(permit);
-    });
-
-    let stream = tokio_stream_from(rx);
-    Response::builder()
-        .status(StatusCode::OK)
-        .header(header::CONTENT_TYPE, "application/x-ndjson")
-        .header(header::CACHE_CONTROL, "no-store")
-        // Some reverse proxies buffer streamed responses unless told not to
-        .header("x-accel-buffering", "no")
-        .body(Body::from_stream(stream))
-        .expect("valid response")
-}
-
-fn tokio_stream_from<T: Send + 'static>(
-    mut rx: tokio::sync::mpsc::Receiver<T>,
-) -> impl futures_util::Stream<Item = T> + Send + 'static {
-    futures_util::stream::poll_fn(move |cx| rx.poll_recv(cx))
-}
@@ -0,0 +1,87 @@
+use axum::body::{Body, to_bytes};
+use axum::extract::Request;
+use axum::http::{HeaderValue, StatusCode, header};
+use axum::middleware::{self, Next};
+use axum::response::Response;
+use std::path::PathBuf;
+use tower::ServiceExt;
+use yaak_web::{Config, router};
+
+fn config(serve: Option<PathBuf>) -> Config {
+    Config {
+        bind: "127.0.0.1:0".parse().unwrap(),
+        serve,
+        allow_private_networks: false,
+        allowed_origins: vec!["*".into()],
+        max_request_bytes: 1024,
+        max_response_bytes: 2048,
+        max_timeout_secs: 5,
+        rate_limit_per_minute: 10,
+        max_concurrent: 2,
+        trust_forwarded_for: false,
+    }
+}
+
+async fn mark_response(request: Request, next: Next) -> Response {
+    let mut response = next.run(request).await;
+    response.headers_mut().insert("x-hosted-middleware", HeaderValue::from_static("yes"));
+    response
+}
+
+#[tokio::test]
+async fn embedded_router_serves_api_without_static_files() {
+    let app = router(config(None)).layer(middleware::from_fn(mark_response));
+    let response = app
+        .clone()
+        .oneshot(Request::builder().uri("/v1/health").body(Body::empty()).unwrap())
+        .await
+        .unwrap();
+    assert_eq!(response.status(), StatusCode::OK);
+    assert_eq!(response.headers()["x-hosted-middleware"], "yes");
+    let body = to_bytes(response.into_body(), 4096).await.unwrap();
+    let health: serde_json::Value = serde_json::from_slice(&body).unwrap();
+    assert_eq!(health["ok"], true);
+    assert_eq!(health["maxResponseBytes"], 2048);
+    assert_eq!(health["maxTimeoutSecs"], 5);
+
+    let response = app.oneshot(Request::new(Body::empty())).await.unwrap();
+    assert_eq!(response.status(), StatusCode::NOT_FOUND);
+}
+
+#[tokio::test]
+async fn middleware_wraps_html_fallback_and_assets_alongside_api() {
+    let dir = tempfile::tempdir().unwrap();
+    std::fs::write(dir.path().join("index.html"), "<html>Yaak</html>").unwrap();
+    std::fs::create_dir(dir.path().join("assets")).unwrap();
+    std::fs::write(dir.path().join("assets/app-abc.js"), "// app").unwrap();
+    let app = router(config(Some(dir.path().into()))).layer(middleware::from_fn(mark_response));
+
+    for path in [
+        "/",
+        "/workspace/example",
+        "/assets/app-abc.js",
+        "/v1/health",
+    ] {
+        let response = app
+            .clone()
+            .oneshot(Request::builder().uri(path).body(Body::empty()).unwrap())
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::OK, "{path}");
+        assert_eq!(response.headers()["x-hosted-middleware"], "yes", "{path}");
+        if path.starts_with("/assets/") {
+            assert_eq!(
+                response.headers()[header::CACHE_CONTROL],
+                "public, max-age=31536000, immutable"
+            );
+        } else if path == "/v1/health" {
+            assert_eq!(response.headers()[header::CONTENT_TYPE], "application/json");
+        } else {
+            assert_eq!(response.headers()[header::CACHE_CONTROL], "no-cache");
+            assert!(
+                response.headers()[header::CONTENT_TYPE].to_str().unwrap().starts_with("text/html")
+            );
+            assert_eq!(to_bytes(response.into_body(), 4096).await.unwrap(), "<html>Yaak</html>");
+        }
+    }
+}