hehe_core/context/
mod.rs

1use crate::types::{AgentId, RequestId, SessionId, Timestamp};
2use std::collections::HashMap;
3use std::sync::Arc;
4use std::time::Duration;
5use tokio_util::sync::CancellationToken;
6
7pub type TraceId = String;
8
9#[derive(Clone)]
10pub struct Context {
11    pub request_id: RequestId,
12    pub trace_id: Option<TraceId>,
13    pub parent_span_id: Option<String>,
14    pub agent_id: Option<AgentId>,
15    pub session_id: Option<SessionId>,
16    pub started_at: Timestamp,
17    pub deadline: Option<Timestamp>,
18    cancellation: CancellationToken,
19    extensions: Arc<HashMap<String, String>>,
20}
21
22impl Context {
23    pub fn new() -> Self {
24        Self {
25            request_id: RequestId::new(),
26            trace_id: None,
27            parent_span_id: None,
28            agent_id: None,
29            session_id: None,
30            started_at: Timestamp::now(),
31            deadline: None,
32            cancellation: CancellationToken::new(),
33            extensions: Arc::new(HashMap::new()),
34        }
35    }
36
37    pub fn with_timeout(mut self, timeout: Duration) -> Self {
38        let deadline_ms = self.started_at.unix_millis() + timeout.as_millis() as i64;
39        self.deadline = Timestamp::from_unix_millis(deadline_ms);
40        self
41    }
42
43    pub fn with_deadline(mut self, deadline: Timestamp) -> Self {
44        self.deadline = Some(deadline);
45        self
46    }
47
48    pub fn with_trace_id(mut self, trace_id: impl Into<String>) -> Self {
49        self.trace_id = Some(trace_id.into());
50        self
51    }
52
53    pub fn with_agent(mut self, agent_id: AgentId) -> Self {
54        self.agent_id = Some(agent_id);
55        self
56    }
57
58    pub fn with_session(mut self, session_id: SessionId) -> Self {
59        self.session_id = Some(session_id);
60        self
61    }
62
63    pub fn with_cancellation(mut self, token: CancellationToken) -> Self {
64        self.cancellation = token;
65        self
66    }
67
68    pub fn child(&self) -> Self {
69        Self {
70            request_id: RequestId::new(),
71            trace_id: self.trace_id.clone(),
72            parent_span_id: Some(self.request_id.to_string()),
73            agent_id: self.agent_id,
74            session_id: self.session_id,
75            started_at: Timestamp::now(),
76            deadline: self.deadline,
77            cancellation: self.cancellation.child_token(),
78            extensions: Arc::clone(&self.extensions),
79        }
80    }
81
82    pub fn is_cancelled(&self) -> bool {
83        self.cancellation.is_cancelled()
84    }
85
86    pub fn cancel(&self) {
87        self.cancellation.cancel()
88    }
89
90    pub fn cancellation_token(&self) -> CancellationToken {
91        self.cancellation.clone()
92    }
93
94    pub fn is_timeout(&self) -> bool {
95        if let Some(deadline) = self.deadline {
96            Timestamp::now() > deadline
97        } else {
98            false
99        }
100    }
101
102    pub fn is_done(&self) -> bool {
103        self.is_cancelled() || self.is_timeout()
104    }
105
106    pub fn remaining(&self) -> Option<Duration> {
107        self.deadline.map(|d| {
108            let now = Timestamp::now().unix_millis();
109            let deadline = d.unix_millis();
110            if deadline > now {
111                Duration::from_millis((deadline - now) as u64)
112            } else {
113                Duration::ZERO
114            }
115        })
116    }
117
118    pub fn elapsed(&self) -> Duration {
119        self.started_at.elapsed()
120    }
121
122    pub fn get_extension(&self, key: &str) -> Option<&String> {
123        self.extensions.get(key)
124    }
125}
126
127impl Default for Context {
128    fn default() -> Self {
129        Self::new()
130    }
131}
132
133impl std::fmt::Debug for Context {
134    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
135        f.debug_struct("Context")
136            .field("request_id", &self.request_id)
137            .field("trace_id", &self.trace_id)
138            .field("agent_id", &self.agent_id)
139            .field("session_id", &self.session_id)
140            .field("started_at", &self.started_at)
141            .field("deadline", &self.deadline)
142            .field("is_cancelled", &self.is_cancelled())
143            .finish()
144    }
145}
hehe_core/context/mod.rs

hehe_core/context/
mod.rs
