descry_tool_core/

context.rs

//! Tool execution context
//!
//! Provides thread-safe, async-friendly context management using DashMap.

use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::any::{Any, TypeId};
use std::sync::Arc;

use crate::JsonObject;

/// Request metadata
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Meta {
    /// Unique request identifier
    pub request_id: Option<String>,
    /// Unix timestamp in milliseconds
    pub timestamp: Option<u64>,
    /// Source of the request (e.g., "cli", "http", "mcp")
    pub source: Option<String>,
}

impl Meta {
    /// Create empty metadata
    pub fn new() -> Self {
        Self::default()
    }

    /// Set request ID
    pub fn with_request_id(mut self, id: impl Into<String>) -> Self {
        self.request_id = Some(id.into());
        self
    }

    /// Set timestamp
    pub fn with_timestamp(mut self, ts: u64) -> Self {
        self.timestamp = Some(ts);
        self
    }

    /// Set source
    pub fn with_source(mut self, source: impl Into<String>) -> Self {
        self.source = Some(source.into());
        self
    }
}

/// Tool execution context
///
/// Thread-safe context using `Arc<DashMap>` for extensions.
/// Can be freely cloned and passed across await points.
///
/// # Thread Safety
///
/// All operations are thread-safe:
/// - `insert`, `get`, `remove` use DashMap's concurrent HashMap
/// - Extensions wrapped in `Arc<T>` for cheap cloning
///
/// # Examples
///
/// ```
/// use descry_tool_core::ToolContext;
/// use std::sync::Arc;
///
/// let ctx = Arc::new(ToolContext::new());
///
/// // Insert extension
/// #[derive(Debug)]
/// struct MyService;
/// ctx.insert(MyService);
/// ```
#[derive(Debug, Clone)]
pub struct ToolContext {
    /// Request metadata
    pub meta: Meta,
    /// Raw arguments (consumed on first access)
    args: Option<JsonObject>,
    /// Thread-safe extensions (DashMap + Arc)
    extensions: Arc<DashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
}

impl ToolContext {
    /// Create a new context with empty args
    pub fn new() -> Self {
        Self {
            meta: Meta::default(),
            args: Some(JsonObject::new()),
            extensions: Arc::new(DashMap::new()),
        }
    }

    /// Create context with args
    pub fn with_args(args: JsonObject) -> Self {
        Self {
            meta: Meta::default(),
            args: Some(args),
            extensions: Arc::new(DashMap::new()),
        }
    }

    /// Create empty context
    pub fn empty() -> Self {
        Self {
            meta: Meta::default(),
            args: None,
            extensions: Arc::new(DashMap::new()),
        }
    }

    /// Set metadata
    pub fn with_meta(mut self, meta: Meta) -> Self {
        self.meta = meta;
        self
    }

    /// Take and consume the raw arguments
    ///
    /// This can only be called once. Subsequent calls will return None.
    pub fn take_args(&mut self) -> Option<JsonObject> {
        self.args.take()
    }

    /// Peek at the raw arguments without consuming
    pub fn peek_args(&self) -> Option<&JsonObject> {
        self.args.as_ref()
    }

    /// Insert an extension into the context
    ///
    /// Returns the previous value if it existed.
    pub fn insert<T>(&self, ext: T) -> Option<Arc<T>>
    where
        T: Send + Sync + 'static,
    {
        self.extensions
            .insert(TypeId::of::<T>(), Arc::new(ext))
            .and_then(|arc| arc.downcast().ok())
    }

    /// Get a reference to an extension
    pub fn get<T>(&self) -> Option<Arc<T>>
    where
        T: Send + Sync + 'static,
    {
        self.extensions
            .get(&TypeId::of::<T>())
            .and_then(|entry| Arc::downcast(entry.clone()).ok())
    }

    /// Remove an extension from the context
    pub fn remove<T>(&self) -> Option<Arc<T>>
    where
        T: Send + Sync + 'static,
    {
        self.extensions
            .remove(&TypeId::of::<T>())
            .and_then(|(_, arc)| Arc::downcast(arc).ok())
    }

    /// Check if an extension exists
    pub fn contains<T>(&self) -> bool
    where
        T: Send + Sync + 'static,
    {
        self.extensions.contains_key(&TypeId::of::<T>())
    }

    /// Clear all extensions
    pub fn clear(&self) {
        self.extensions.clear();
    }

    /// Get extension count
    pub fn len(&self) -> usize {
        self.extensions.len()
    }

    /// Check if extensions is empty
    pub fn is_empty(&self) -> bool {
        self.extensions.is_empty()
    }
}

impl Default for ToolContext {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_meta_builder() {
        let meta = Meta::new()
            .with_request_id("test-123")
            .with_timestamp(1234567890)
            .with_source("test");

        assert_eq!(meta.request_id, Some("test-123".to_string()));
        assert_eq!(meta.timestamp, Some(1234567890));
        assert_eq!(meta.source, Some("test".to_string()));
    }

    #[test]
    fn test_context_new() {
        let ctx = ToolContext::new();
        assert!(ctx.peek_args().is_some());
        assert_eq!(ctx.meta.request_id, None);
    }

    #[test]
    fn test_context_take_args() {
        let mut ctx = ToolContext::with_args(
            serde_json::json!({"key": "value"})
                .as_object()
                .unwrap()
                .clone(),
        );

        let taken = ctx.take_args().unwrap();
        assert!(taken.contains_key("key"));

        // Second call should return None
        assert!(ctx.take_args().is_none());
    }

    #[test]
    fn test_context_extensions() {
        let ctx = ToolContext::new();

        #[derive(Debug, PartialEq)]
        struct MyExt {
            value: i32,
        }

        // Insert (returns None for new key)
        assert!(ctx.insert(MyExt { value: 42 }).is_none());

        // Get
        let ext = ctx.get::<MyExt>().unwrap();
        assert_eq!(ext.value, 42);

        // Contains
        assert!(ctx.contains::<MyExt>());

        // Remove
        let removed = ctx.remove::<MyExt>().unwrap();
        assert_eq!(removed.value, 42);
        assert!(!ctx.contains::<MyExt>());
    }

    #[test]
    fn test_context_clone() {
        #[derive(Debug, PartialEq)]
        struct MyExt {
            value: i32,
        }

        let ctx = ToolContext::new();
        ctx.insert(MyExt { value: 42 });

        // Clone should preserve extensions
        let cloned = ctx.clone();
        let ext = cloned.get::<MyExt>().unwrap();
        assert_eq!(ext.value, 42);
    }

    #[test]
    fn test_context_thread_safety() {
        use std::sync::Arc;
        use std::thread;

        let ctx = Arc::new(ToolContext::new());
        let ctx_clone = ctx.clone();

        #[derive(Debug)]
        struct Counter;

        // Insert in main thread
        ctx.insert(Counter);

        // Access in spawned thread
        let handle = thread::spawn(move || {
            assert!(ctx_clone.contains::<Counter>());
        });

        handle.join().unwrap();
        assert!(ctx.contains::<Counter>());
    }
}