Skip to main content

subcog/mcp/
session.rs

1//! Session state tracking for MCP clients.
2//!
3//! Tracks whether the current session has been properly initialized via `subcog_init`.
4//! Provides lightweight hints to uninitiated sessions for the first few tool calls.
5
6use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
7
8/// Maximum number of hints to show for uninitialized sessions.
9const MAX_HINTS: u32 = 3;
10
11/// Global session state tracker.
12///
13/// Uses atomic operations for thread-safe state management without locks.
14/// State is reset on server restart (ephemeral by design).
15static INITIALIZED: AtomicBool = AtomicBool::new(false);
16static HINT_COUNT: AtomicU32 = AtomicU32::new(0);
17
18/// Marks the session as initialized.
19///
20/// Called by `subcog_init` after successful initialization.
21pub fn mark_initialized() {
22    INITIALIZED.store(true, Ordering::Release);
23}
24
25/// Checks if the session has been initialized.
26#[must_use]
27pub fn is_initialized() -> bool {
28    INITIALIZED.load(Ordering::Acquire)
29}
30
31/// Checks if a hint should be shown and increments the counter.
32///
33/// Returns `true` if:
34/// - Session is not initialized AND
35/// - Hint count is below `MAX_HINTS`
36///
37/// Atomically increments the hint counter when returning `true`.
38#[must_use]
39pub fn should_show_hint() -> bool {
40    if is_initialized() {
41        return false;
42    }
43
44    // Atomically check and increment hint count
45    let current = HINT_COUNT.fetch_add(1, Ordering::AcqRel);
46    current < MAX_HINTS
47}
48
49/// Returns the initialization hint message.
50///
51/// This is appended to tool responses when `should_show_hint()` returns `true`.
52#[must_use]
53pub const fn get_hint_message() -> &'static str {
54    "\n\n---\n💡 **Tip**: Call `subcog_init` at session start to load memory context and best practices."
55}
56
57/// Resets session state (primarily for testing).
58#[cfg(test)]
59pub fn reset() {
60    INITIALIZED.store(false, Ordering::Release);
61    HINT_COUNT.store(0, Ordering::Release);
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    #[test]
69    fn test_initialization_flow() {
70        reset();
71
72        // Initially not initialized
73        assert!(!is_initialized());
74
75        // Mark as initialized
76        mark_initialized();
77        assert!(is_initialized());
78    }
79
80    #[test]
81    fn test_hint_limiting() {
82        reset();
83
84        // Should show hints for first MAX_HINTS calls
85        for i in 0..MAX_HINTS {
86            assert!(should_show_hint(), "Should show hint for call {i}");
87        }
88
89        // Should stop showing hints after MAX_HINTS
90        assert!(!should_show_hint(), "Should not show hint after MAX_HINTS");
91    }
92
93    #[test]
94    fn test_no_hints_when_initialized() {
95        reset();
96
97        // Initialize first
98        mark_initialized();
99
100        // Should never show hints when initialized
101        for _ in 0..10 {
102            assert!(
103                !should_show_hint(),
104                "Should not show hints when initialized"
105            );
106        }
107    }
108
109    #[test]
110    fn test_hint_message_content() {
111        let msg = get_hint_message();
112        assert!(msg.contains("subcog_init"));
113        assert!(msg.contains("session start"));
114    }
115}