Genarrative/server-rs/crates/api-server/src/modules/editor_project.rs

use axum::{
    Router,
    extract::{DefaultBodyLimit, Request},
    handler::Handler,
    http::{HeaderValue, header},
    middleware,
    middleware::Next,
    response::Response,
    routing::{get, patch, post},
};

use crate::{
    auth::require_bearer_auth,
    editor_agent::api::editor_agent_message,
    editor_agent::{
        cancel_editor_agent_tool_call, confirm_editor_agent_tool_call,
        create_editor_agent_conversation, delete_editor_agent_conversation,
        get_editor_agent_conversation, list_editor_agent_conversations,
    },
    editor_project::{
        EDITOR_LAYOUT_REQUEST_BODY_MAX_BYTES, create_editor_asset, create_editor_asset_folder,
        create_editor_project, create_editor_project_resource, delete_editor_asset,
        delete_editor_asset_folder, delete_editor_project, edit_editor_image,
        extract_editor_ui_design_assets, generate_editor_image, generate_editor_scene,
        get_editor_asset_library, get_editor_generation_pricing, get_editor_project,
        list_editor_projects, list_public_editor_project_resources, load_recent_editor_project,
        remove_editor_image_background, rename_editor_project, save_editor_project_layout,
        snap_editor_image_to_pixel_art, submit_editor_asset_showcase,
        toggle_editor_showcase_asset_like, update_editor_asset, update_editor_asset_folder,
        update_editor_project_resource_showcase,
    },
    editor_project_icon::{
        generate_editor_icon_spritesheet, generate_icon_spec, split_editor_icon_spritesheet,
    },
    llm::icon_specs::{refine_art_style, refine_game_play},
    state::AppState,
};

const EDITOR_AGENT_MESSAGE_BODY_LIMIT_BYTES: usize = 2 * 1024 * 1024;

async fn add_showcase_private_response_headers(request: Request, next: Next) -> Response {
    let mut response = next.run(request).await;
    response.headers_mut().insert(
        header::CACHE_CONTROL,
        HeaderValue::from_static("private, no-store"),
    );
    response
        .headers_mut()
        .append(header::VARY, HeaderValue::from_static("Authorization"));
    response
}

pub fn router(state: AppState) -> Router<AppState> {
    Router::new()
        .route(
            "/api/editor/generation-pricing",
            get(get_editor_generation_pricing),
        )
        .route(
            "/api/editor/showcase/resources",
            get(list_public_editor_project_resources)
                .route_layer(middleware::from_fn(add_showcase_private_response_headers)),
        )
        .route(
            "/api/editor/projects/recent",
            get(load_recent_editor_project).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/projects",
            get(list_editor_projects)
                .post(create_editor_project)
                .route_layer(middleware::from_fn_with_state(
                    state.clone(),
                    require_bearer_auth,
                )),
        )
        .route(
            "/api/editor/projects/{project_id}",
            get(get_editor_project)
                .patch(
                    save_editor_project_layout
                        .layer(DefaultBodyLimit::max(EDITOR_LAYOUT_REQUEST_BODY_MAX_BYTES)),
                )
                .delete(delete_editor_project)
                .route_layer(middleware::from_fn_with_state(
                    state.clone(),
                    require_bearer_auth,
                )),
        )
        .route(
            "/api/editor/projects/{project_id}/metadata",
            patch(rename_editor_project).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/projects/{project_id}/resources",
            post(create_editor_project_resource).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/projects/{project_id}/agent-conversations",
            get(list_editor_agent_conversations)
                .post(create_editor_agent_conversation)
                .route_layer(middleware::from_fn_with_state(
                    state.clone(),
                    require_bearer_auth,
                )),
        )
        .route(
            "/api/editor/agent-conversations/{conversation_id}",
            get(get_editor_agent_conversation)
                .delete(delete_editor_agent_conversation)
                .route_layer(middleware::from_fn_with_state(
                    state.clone(),
                    require_bearer_auth,
                )),
        )
        .route(
            "/api/editor/agent-conversations/{conversation_id}/messages",
            post(editor_agent_message)
                .layer(DefaultBodyLimit::max(EDITOR_AGENT_MESSAGE_BODY_LIMIT_BYTES))
                .route_layer(middleware::from_fn_with_state(
                    state.clone(),
                    require_bearer_auth,
                )),
        )
        .route(
            "/api/editor/agent-conversations/{conversation_id}/messages/{message_id}/confirm",
            post(confirm_editor_agent_tool_call).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/agent-conversations/{conversation_id}/messages/{message_id}/cancel",
            post(cancel_editor_agent_tool_call).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/project-resources/{resource_id}/showcase",
            patch(update_editor_project_resource_showcase).route_layer(
                middleware::from_fn_with_state(state.clone(), require_bearer_auth),
            ),
        )
        .route(
            "/api/editor/assets/library",
            get(get_editor_asset_library).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/assets/folders",
            post(create_editor_asset_folder).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/assets/folders/{folder_id}",
            patch(update_editor_asset_folder)
                .delete(delete_editor_asset_folder)
                .route_layer(middleware::from_fn_with_state(
                    state.clone(),
                    require_bearer_auth,
                )),
        )
        .route(
            "/api/editor/assets",
            post(create_editor_asset).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/assets/{asset_id}",
            patch(update_editor_asset)
                .delete(delete_editor_asset)
                .route_layer(middleware::from_fn_with_state(
                    state.clone(),
                    require_bearer_auth,
                )),
        )
        .route(
            "/api/editor/assets/{asset_id}/showcase-submissions",
            post(submit_editor_asset_showcase).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/showcase/assets/{showcase_id}/likes",
            post(toggle_editor_showcase_asset_like)
                .route_layer(middleware::from_fn_with_state(
                    state.clone(),
                    require_bearer_auth,
                ))
                .route_layer(middleware::from_fn(add_showcase_private_response_headers)),
        )
        .route(
            "/api/editor/images/generations",
            post(generate_editor_image).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/scenes/generations",
            post(generate_editor_scene).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/llm/icon-specs/refine-game-play",
            post(refine_game_play).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/llm/icon-specs/refine-art-style",
            post(refine_art_style).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/icon-specs/generations",
            post(generate_icon_spec).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/images/edits",
            post(edit_editor_image).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/images/background-removals",
            post(remove_editor_image_background).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/images/pixel-art-snaps",
            post(snap_editor_image_to_pixel_art).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/icon-spritesheets/generations",
            post(generate_editor_icon_spritesheet).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/icon-spritesheets/slices",
            post(split_editor_icon_spritesheet).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
        .route(
            "/api/editor/ui-designs/assets/extractions",
            post(extract_editor_ui_design_assets).route_layer(middleware::from_fn_with_state(
                state.clone(),
                require_bearer_auth,
            )),
        )
}

#[cfg(test)]
mod tests {
    use axum::{
        Router,
        body::Body,
        http::{Request, header},
        middleware,
        response::{IntoResponse, Response},
        routing::get,
    };
    use tower::ServiceExt;

    use super::add_showcase_private_response_headers;

    async fn response_with_existing_vary() -> Response {
        let mut response = "ok".into_response();
        response
            .headers_mut()
            .insert(header::VARY, "Origin".parse().expect("valid Vary value"));
        response
    }

    #[tokio::test]
    async fn showcase_private_headers_preserve_existing_vary_fields() {
        let app = Router::new().route(
            "/",
            get(response_with_existing_vary)
                .route_layer(middleware::from_fn(add_showcase_private_response_headers)),
        );

        let response = app
            .oneshot(
                Request::builder()
                    .uri("/")
                    .body(Body::empty())
                    .expect("request should build"),
            )
            .await
            .expect("request should succeed");
        let vary_values = response
            .headers()
            .get_all(header::VARY)
            .iter()
            .map(|value| value.to_str().expect("Vary should be ASCII"))
            .collect::<Vec<_>>();

        assert_eq!(vary_values, vec!["Origin", "Authorization"]);
    }
}