use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;
use turul_mcp_session_storage::{BoxedSessionStorage, SessionView};
pub struct StorageBackedSessionView {
session_id: String,
storage: Arc<BoxedSessionStorage>,
}
impl StorageBackedSessionView {
pub fn new(session_id: String, storage: Arc<BoxedSessionStorage>) -> Self {
Self {
session_id,
storage,
}
}
}
#[async_trait]
impl SessionView for StorageBackedSessionView {
fn session_id(&self) -> &str {
&self.session_id
}
async fn get_state(&self, key: &str) -> Result<Option<Value>, String> {
let session = self
.storage
.get_session(&self.session_id)
.await
.map_err(|e| format!("Failed to get session: {}", e))?;
Ok(session.and_then(|s| s.state.get(key).cloned()))
}
async fn set_state(&self, key: &str, value: Value) -> Result<(), String> {
let mut session = self
.storage
.get_session(&self.session_id)
.await
.map_err(|e| format!("Failed to get session: {}", e))?
.ok_or_else(|| format!("Session '{}' not found", self.session_id))?;
session.state.insert(key.to_string(), value);
session.last_activity = chrono::Utc::now().timestamp_millis() as u64;
self.storage
.update_session(session)
.await
.map_err(|e| format!("Failed to update session: {}", e))
}
async fn get_metadata(&self, key: &str) -> Result<Option<Value>, String> {
let prefixed_key = format!("__meta__:{}", key);
let session = self
.storage
.get_session(&self.session_id)
.await
.map_err(|e| format!("Failed to get session: {}", e))?;
Ok(session.and_then(|s| s.metadata.get(&prefixed_key).cloned()))
}
async fn set_metadata(&self, key: &str, value: Value) -> Result<(), String> {
let prefixed_key = format!("__meta__:{}", key);
let mut session = self
.storage
.get_session(&self.session_id)
.await
.map_err(|e| format!("Failed to get session: {}", e))?
.ok_or_else(|| format!("Session '{}' not found", self.session_id))?;
session.metadata.insert(prefixed_key, value);
session.last_activity = chrono::Utc::now().timestamp_millis() as u64;
self.storage
.update_session(session)
.await
.map_err(|e| format!("Failed to update session: {}", e))
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use turul_mcp_protocol::ServerCapabilities;
use turul_mcp_session_storage::{BoxedSessionStorage, InMemorySessionStorage};
#[tokio::test]
async fn test_session_view_state() {
let storage: Arc<BoxedSessionStorage> = Arc::new(InMemorySessionStorage::new());
let session_info = storage
.create_session(ServerCapabilities::default())
.await
.unwrap();
let session_id = session_info.session_id.clone();
let view = StorageBackedSessionView::new(session_id.clone(), Arc::clone(&storage));
assert_eq!(view.get_state("key1").await.unwrap(), None);
view.set_state("key1", json!("value1")).await.unwrap();
assert_eq!(view.get_state("key1").await.unwrap(), Some(json!("value1")));
view.set_state("key2", json!({"nested": "object"}))
.await
.unwrap();
assert_eq!(
view.get_state("key2").await.unwrap(),
Some(json!({"nested": "object"}))
);
}
#[tokio::test]
async fn test_session_view_metadata() {
let storage: Arc<BoxedSessionStorage> = Arc::new(InMemorySessionStorage::new());
let session_info = storage
.create_session(ServerCapabilities::default())
.await
.unwrap();
let session_id = session_info.session_id.clone();
let view = StorageBackedSessionView::new(session_id.clone(), Arc::clone(&storage));
assert_eq!(view.get_metadata("meta1").await.unwrap(), None);
view.set_metadata("meta1", json!("metadata_value"))
.await
.unwrap();
assert_eq!(
view.get_metadata("meta1").await.unwrap(),
Some(json!("metadata_value"))
);
let session = storage.get_session(&session_id).await.unwrap().unwrap();
assert_eq!(
session.metadata.get("__meta__:meta1"),
Some(&json!("metadata_value"))
);
}
#[tokio::test]
async fn test_session_view_session_id() {
let storage: Arc<BoxedSessionStorage> = Arc::new(InMemorySessionStorage::new());
let session_info = storage
.create_session(ServerCapabilities::default())
.await
.unwrap();
let view =
StorageBackedSessionView::new(session_info.session_id.clone(), Arc::clone(&storage));
assert_eq!(view.session_id(), &session_info.session_id);
}
#[tokio::test]
async fn test_session_view_nonexistent_session() {
let storage: Arc<BoxedSessionStorage> = Arc::new(InMemorySessionStorage::new());
let view = StorageBackedSessionView::new("nonexistent".to_string(), Arc::clone(&storage));
assert_eq!(view.get_state("key").await.unwrap(), None);
let result = view.set_state("key", json!("value")).await;
assert!(result.is_err());
assert!(result.unwrap_err().contains("not found"));
}
}