rust_mcp_sdk/mcp_traits/
request_id_gen.rs

1use std::sync::atomic::AtomicI64;
2
3use crate::schema::{schema_utils::McpMessage, RequestId};
4use async_trait::async_trait;
5
6/// A trait for generating and managing request IDs in a thread-safe manner.
7///
8/// Implementors provide functionality to generate unique request IDs, retrieve the last
9/// generated ID, and reset the ID counter.
10#[async_trait]
11pub trait RequestIdGen: Send + Sync {
12    fn next_request_id(&self) -> RequestId;
13    #[allow(unused)]
14    fn last_request_id(&self) -> Option<RequestId>;
15    #[allow(unused)]
16    fn reset_to(&self, id: u64);
17
18    /// Determines the request ID for an outgoing MCP message.
19    ///
20    /// For requests, generates a new ID using the internal counter. For responses or errors,
21    /// uses the provided `request_id`. Notifications receive no ID.
22    ///
23    /// # Arguments
24    /// * `message` - The MCP message to evaluate.
25    /// * `request_id` - An optional existing request ID (required for responses/errors).
26    ///
27    /// # Returns
28    /// An `Option<RequestId>`: `Some` for requests or responses/errors, `None` for notifications.
29    fn request_id_for_message(
30        &self,
31        message: &dyn McpMessage,
32        request_id: Option<RequestId>,
33    ) -> Option<RequestId> {
34        // we need to produce next request_id for requests
35        if message.is_request() {
36            // request_id should be None for requests
37            assert!(request_id.is_none());
38            Some(self.next_request_id())
39        } else if !message.is_notification() {
40            // `request_id` must not be `None` for errors, notifications and responses
41            assert!(request_id.is_some());
42            request_id
43        } else {
44            None
45        }
46    }
47}
48
49pub struct RequestIdGenNumeric {
50    message_id_counter: AtomicI64,
51    last_message_id: AtomicI64,
52}
53
54impl RequestIdGenNumeric {
55    pub fn new(initial_id: Option<u64>) -> Self {
56        Self {
57            message_id_counter: AtomicI64::new(initial_id.unwrap_or(0) as i64),
58            last_message_id: AtomicI64::new(-1),
59        }
60    }
61}
62
63impl RequestIdGen for RequestIdGenNumeric {
64    /// Generates the next unique request ID as an integer.
65    ///
66    /// Increments the internal counter atomically and updates the last generated ID.
67    /// Uses `Relaxed` ordering for performance, as the counter only needs to ensure unique IDs.
68    fn next_request_id(&self) -> RequestId {
69        let id = self
70            .message_id_counter
71            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
72        // Store the new ID as the last generated ID
73        self.last_message_id
74            .store(id, std::sync::atomic::Ordering::Relaxed);
75        RequestId::Integer(id)
76    }
77
78    /// Returns the last generated request ID, if any.
79    ///
80    /// Returns `None` if no ID has been generated (indicated by a sentinel value of -1).
81    /// Uses `Relaxed` ordering since the read operation doesn’t require synchronization
82    /// with other memory operations beyond atomicity.
83    fn last_request_id(&self) -> Option<RequestId> {
84        let last_id = self
85            .last_message_id
86            .load(std::sync::atomic::Ordering::Relaxed);
87        if last_id == -1 {
88            None
89        } else {
90            Some(RequestId::Integer(last_id))
91        }
92    }
93
94    /// Resets the internal counter to the specified ID.
95    ///
96    /// The provided `id` (u64) is converted to i64 and stored atomically.
97    fn reset_to(&self, id: u64) {
98        self.message_id_counter
99            .store(id as i64, std::sync::atomic::Ordering::Relaxed);
100    }
101}
rust_mcp_sdk/mcp_traits/request_id_gen.rs

rust_mcp_sdk/mcp_traits/
request_id_gen.rs
