mnote/rust/crates/mnote-web/src/context.rs

use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, Uri};
use serde::Serialize;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};

static REQUEST_COUNTER: AtomicU64 = AtomicU64::new(1);

const HEADER_REQUEST_ID: &str = "x-request-id";
const HEADER_TRACE_ID: &str = "x-trace-id";
const HEADER_TENANT_ID: &str = "x-mnote-tenant-id";
const HEADER_WORKSPACE_ID: &str = "x-mnote-workspace-id";
const HEADER_DEPLOYMENT_ID: &str = "x-mnote-deployment-id";
const HEADER_PROJECT_ID: &str = "x-mnote-project-id";
const HEADER_ACTOR_ID: &str = "x-mnote-actor-id";
const HEADER_ACTOR_TYPE: &str = "x-mnote-actor-type";
const HEADER_SESSION_ID: &str = "x-mnote-session-id";
const HEADER_SOURCE_CHANNEL: &str = "x-mnote-source-channel";
const HEADER_SOURCE_CLIENT: &str = "x-mnote-source-client";
const HEADER_IDEMPOTENCY_KEY: &str = "x-idempotency-key";

#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct TraceContext {
    pub request_id: String,
    pub trace_id: String,
    pub method: String,
    pub path: String,
}

#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct AuthContext {
    pub authorization: Option<String>,
    pub actor_id: String,
    pub actor_type: String,
    pub session_id: Option<String>,
}

#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct WorkspaceContext {
    pub workspace_id: Option<String>,
    pub tenant_id: Option<String>,
    pub deployment_id: Option<String>,
    pub project_id: Option<String>,
}

#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct SourceContext {
    pub channel: String,
    pub client: String,
    pub idempotency_key: Option<String>,
}

#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct RequestContext {
    pub trace: TraceContext,
    pub auth: AuthContext,
    pub workspace: WorkspaceContext,
    pub source: SourceContext,
}

impl RequestContext {
    pub fn from_http_parts(method: &Method, uri: &Uri, headers: &HeaderMap) -> Self {
        let request_id = header_or_generated(headers, HEADER_REQUEST_ID, "req");
        let trace_id = header_or_generated(headers, HEADER_TRACE_ID, "trace");

        Self {
            trace: TraceContext {
                request_id,
                trace_id,
                method: method.as_str().to_string(),
                path: uri.path().to_string(),
            },
            auth: AuthContext {
                authorization: header_value(headers, axum::http::header::AUTHORIZATION.as_str()),
                actor_id: header_value(headers, HEADER_ACTOR_ID)
                    .unwrap_or_else(|| "anonymous".into()),
                actor_type: header_value(headers, HEADER_ACTOR_TYPE)
                    .unwrap_or_else(|| "anonymous".into()),
                session_id: header_value(headers, HEADER_SESSION_ID),
            },
            workspace: WorkspaceContext {
                workspace_id: header_value(headers, HEADER_WORKSPACE_ID),
                tenant_id: header_value(headers, HEADER_TENANT_ID),
                deployment_id: header_value(headers, HEADER_DEPLOYMENT_ID),
                project_id: header_value(headers, HEADER_PROJECT_ID),
            },
            source: SourceContext {
                channel: header_value(headers, HEADER_SOURCE_CHANNEL)
                    .unwrap_or_else(|| "http".into()),
                client: header_value(headers, HEADER_SOURCE_CLIENT)
                    .unwrap_or_else(|| "mnote-web".into()),
                idempotency_key: header_value(headers, HEADER_IDEMPOTENCY_KEY),
            },
        }
    }

    pub fn apply_response_headers(&self, headers: &mut HeaderMap) {
        insert_header(headers, HEADER_REQUEST_ID, &self.trace.request_id);
        insert_header(headers, HEADER_TRACE_ID, &self.trace.trace_id);
        if let Some(workspace_id) = &self.workspace.workspace_id {
            insert_header(headers, HEADER_WORKSPACE_ID, workspace_id);
        }
    }
}

fn header_or_generated(headers: &HeaderMap, key: &str, prefix: &str) -> String {
    header_value(headers, key).unwrap_or_else(|| generate_id(prefix))
}

fn header_value(headers: &HeaderMap, key: &str) -> Option<String> {
    headers
        .get(key)
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToOwned::to_owned)
}

fn generate_id(prefix: &str) -> String {
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis();
    let counter = REQUEST_COUNTER.fetch_add(1, Ordering::Relaxed);
    format!("{prefix}_{now}_{counter}")
}

fn insert_header(headers: &mut HeaderMap, key: &str, value: &str) {
    let Ok(name) = HeaderName::from_lowercase(key.as_bytes()) else {
        return;
    };
    let Ok(value) = HeaderValue::from_str(value) else {
        return;
    };
    headers.insert(name, value);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn request_context_uses_headers_when_present() {
        let mut headers = HeaderMap::new();
        headers.insert(HEADER_REQUEST_ID, HeaderValue::from_static("req_demo"));
        headers.insert(HEADER_TRACE_ID, HeaderValue::from_static("trace_demo"));
        headers.insert(HEADER_WORKSPACE_ID, HeaderValue::from_static("ws_demo"));
        headers.insert(HEADER_ACTOR_ID, HeaderValue::from_static("user_demo"));

        let context = RequestContext::from_http_parts(
            &Method::POST,
            &"/api/hermes/bridge".parse::<Uri>().expect("uri"),
            &headers,
        );

        assert_eq!(context.trace.request_id, "req_demo");
        assert_eq!(context.trace.trace_id, "trace_demo");
        assert_eq!(context.workspace.workspace_id.as_deref(), Some("ws_demo"));
        assert_eq!(context.auth.actor_id, "user_demo");
    }
}