github_mcp/core/
correlation_context.rs

1// GitHub v3 REST API MCP server — generated by mcpify. Do not hand-edit.
2
3use std::future::Future;
4
5use tokio::task_local;
6use uuid::Uuid;
7
8task_local! {
9    static TRACE_ID: String;
10}
11
12/// Runs `fut` with a fresh (or given) correlation ID bound for its entire
13/// async call tree — the `tokio::task_local!` equivalent of
14/// `targets::typescript`'s `AsyncLocalStorage`-based `withCorrelationId`.
15pub async fn with_correlation_id<F, T>(trace_id: Option<String>, fut: F) -> T
16where
17    F: Future<Output = T>,
18{
19    let id = trace_id.unwrap_or_else(|| Uuid::new_v4().to_string());
20    TRACE_ID.scope(id, fut).await
21}
22
23pub fn correlation_id() -> Option<String> {
24    TRACE_ID.try_with(|id| id.clone()).ok()
25}
26
27#[cfg(test)]
28mod tests {
29    use super::*;
30
31    #[tokio::test]
32    async fn binds_a_generated_trace_id_for_the_call_tree() {
33        with_correlation_id(None, async {
34            assert!(correlation_id().is_some());
35        })
36        .await;
37    }
38
39    #[tokio::test]
40    async fn reuses_a_provided_trace_id() {
41        with_correlation_id(Some("fixed-id".to_string()), async {
42            assert_eq!(correlation_id().as_deref(), Some("fixed-id"));
43        })
44        .await;
45    }
46
47    #[tokio::test]
48    async fn reports_none_outside_any_bound_scope() {
49        assert_eq!(correlation_id(), None);
50    }
51}
github_mcp/core/correlation_context.rs

github_mcp/core/
correlation_context.rs
