qwencode_rs/query/
session.rs

1use anyhow::Result;
2use tokio_util::sync::CancellationToken;
3use tracing::{debug, info};
4use uuid::Uuid;
5
6use crate::types::permission::PermissionMode;
7
8/// Handle for an active query session
9pub struct QueryHandle {
10    session_id: String,
11    cancel_token: CancellationToken,
12    is_closed: bool,
13}
14
15impl QueryHandle {
16    /// Create a new query handle with a session ID
17    pub fn new(session_id: Option<String>) -> Self {
18        let session_id = session_id.unwrap_or_else(|| Uuid::new_v4().to_string());
19        debug!("Creating new QueryHandle with session_id: {}", session_id);
20
21        QueryHandle {
22            session_id,
23            cancel_token: CancellationToken::new(),
24            is_closed: false,
25        }
26    }
27
28    /// Get the session ID
29    pub fn session_id(&self) -> &str {
30        &self.session_id
31    }
32
33    /// Check if the session is closed
34    pub fn is_closed(&self) -> bool {
35        self.is_closed
36    }
37
38    /// Get the cancellation token
39    pub fn cancellation_token(&self) -> CancellationToken {
40        self.cancel_token.clone()
41    }
42
43    /// Interrupt the current operation
44    pub async fn interrupt(&self) -> Result<()> {
45        info!("Interrupting session {}", self.session_id);
46        self.cancel_token.cancel();
47        Ok(())
48    }
49
50    /// Set the permission mode (would update the running CLI process)
51    pub async fn set_permission_mode(&self, _mode: PermissionMode) -> Result<()> {
52        // TODO: Implement mode change via CLI communication
53        info!("Setting permission mode for session {}", self.session_id);
54        Ok(())
55    }
56
57    /// Set the model (would update the running CLI process)
58    pub async fn set_model(&self, _model: &str) -> Result<()> {
59        // TODO: Implement model change via CLI communication
60        info!("Setting model for session {}", self.session_id);
61        Ok(())
62    }
63
64    /// Close the session
65    pub async fn close(&mut self) -> Result<()> {
66        if !self.is_closed {
67            info!("Closing session {}", self.session_id);
68            self.cancel_token.cancel();
69            self.is_closed = true;
70        }
71        Ok(())
72    }
73}
74
75impl Drop for QueryHandle {
76    fn drop(&mut self) {
77        if !self.is_closed {
78            debug!("QueryHandle dropped without explicit close, cancelling");
79            self.cancel_token.cancel();
80        }
81    }
82}
83
84/// Generate a session ID
85pub fn generate_session_id() -> String {
86    Uuid::new_v4().to_string()
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    #[test]
94    fn test_query_handle_creation() {
95        let handle = QueryHandle::new(None);
96
97        assert!(!handle.session_id().is_empty());
98        assert!(!handle.is_closed());
99    }
100
101    #[test]
102    fn test_query_handle_with_session_id() {
103        let handle = QueryHandle::new(Some("custom-session-id".to_string()));
104
105        assert_eq!(handle.session_id(), "custom-session-id");
106    }
107
108    #[test]
109    fn test_query_handle_cancellation_token() {
110        let handle = QueryHandle::new(None);
111        let token = handle.cancellation_token();
112
113        assert!(!token.is_cancelled());
114    }
115
116    #[tokio::test]
117    async fn test_query_handle_interrupt() {
118        let handle = QueryHandle::new(None);
119        let token = handle.cancellation_token();
120
121        handle.interrupt().await.unwrap();
122        assert!(token.is_cancelled());
123    }
124
125    #[tokio::test]
126    async fn test_query_handle_close() {
127        let mut handle = QueryHandle::new(None);
128
129        assert!(!handle.is_closed());
130        handle.close().await.unwrap();
131        assert!(handle.is_closed());
132    }
133
134    #[tokio::test]
135    async fn test_query_handle_close_multiple_times() {
136        let mut handle = QueryHandle::new(None);
137
138        handle.close().await.unwrap();
139        handle.close().await.unwrap(); // Should not panic
140
141        assert!(handle.is_closed());
142    }
143
144    #[test]
145    fn test_generate_session_id() {
146        let id1 = generate_session_id();
147        let id2 = generate_session_id();
148
149        assert!(!id1.is_empty());
150        assert!(!id2.is_empty());
151        assert_ne!(id1, id2);
152    }
153
154    #[test]
155    fn test_session_id_format() {
156        let id = generate_session_id();
157
158        // UUID v4 format: 8-4-4-4-12 hexadecimal characters
159        let parts: Vec<&str> = id.split('-').collect();
160        assert_eq!(parts.len(), 5);
161        assert_eq!(parts[0].len(), 8);
162        assert_eq!(parts[1].len(), 4);
163        assert_eq!(parts[2].len(), 4);
164        assert_eq!(parts[3].len(), 4);
165        assert_eq!(parts[4].len(), 12);
166    }
167
168    #[tokio::test]
169    async fn test_query_handle_set_permission_mode() {
170        let handle = QueryHandle::new(None);
171
172        let result = handle.set_permission_mode(PermissionMode::Yolo).await;
173        assert!(result.is_ok());
174    }
175
176    #[tokio::test]
177    async fn test_query_handle_set_model() {
178        let handle = QueryHandle::new(None);
179
180        let result = handle.set_model("qwen-max").await;
181        assert!(result.is_ok());
182    }
183
184    #[test]
185    fn test_query_handle_drop_cancels_token() {
186        let token;
187        {
188            let handle = QueryHandle::new(None);
189            token = handle.cancellation_token();
190            assert!(!token.is_cancelled());
191        } // handle dropped
192
193        // After drop, token should be cancelled
194        assert!(token.is_cancelled());
195    }
196}
qwencode_rs/query/session.rs

qwencode_rs/query/
session.rs
