everruns-core 0.9.0

//! Session Storage Capability
//!
//! This capability provides tools for session-scoped key/value and secret storage.
//! Data persists for the session duration.
//!
//! Tools provided:
//! - `kv_store`: Key/value storage operations (set, get, delete, list)
//! - `secret_store`: Encrypted secret storage operations (set, get, delete, list)

use super::{Capability, CapabilityStatus};
use crate::tool_types::ToolHints;
use crate::tools::{Tool, ToolExecutionResult};
use crate::traits::ToolContext;
use async_trait::async_trait;
use serde_json::{Value, json};

const INTERNAL_SECRET_PREFIXES: &[&str] = &["browserless_internal:"];

fn is_internal_secret_name(name: &str) -> bool {
    INTERNAL_SECRET_PREFIXES
        .iter()
        .any(|prefix| name.starts_with(prefix))
}

/// Session Storage capability - provides key/value and secret storage for sessions
pub struct SessionStorageCapability;

impl Capability for SessionStorageCapability {
    fn id(&self) -> &str {
        "session_storage"
    }

    fn name(&self) -> &str {
        "Storage"
    }

    fn description(&self) -> &str {
        r#"Tools to store and retrieve key/value pairs and encrypted secrets within a session.

> [!NOTE]
> Data persists for the session duration. Secrets are encrypted at rest.

> [!TIP]
> Use key/value storage for general data. Use secrets for sensitive information like API keys or tokens."#
    }

    fn status(&self) -> CapabilityStatus {
        CapabilityStatus::Available
    }

    fn icon(&self) -> Option<&str> {
        Some("database")
    }

    fn category(&self) -> Option<&str> {
        Some("Storage")
    }

    fn system_prompt_addition(&self) -> Option<&str> {
        Some(
            "Use `kv_store` for general data. Use `secret_store` for sensitive data (API keys, tokens, credentials) — secrets are encrypted at rest. Keys are unique per session; storing with the same key overwrites.",
        )
    }

    fn tools(&self) -> Vec<Box<dyn Tool>> {
        vec![Box::new(KvStoreTool), Box::new(SecretStoreTool)]
    }

    fn features(&self) -> Vec<&'static str> {
        vec!["secrets", "key_value"]
    }
}

// ============================================================================
// KvStoreTool - Unified key/value storage tool
// ============================================================================

/// Tool for key/value storage operations
pub struct KvStoreTool;

#[async_trait]
impl Tool for KvStoreTool {
    fn name(&self) -> &str {
        "kv_store"
    }

    fn display_name(&self) -> Option<&str> {
        Some("Key-Value Store")
    }

    fn description(&self) -> &str {
        "Key/value storage operations: set, get, delete, or list keys."
    }

    fn parameters_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "enum": ["set", "get", "delete", "list"],
                    "description": "The operation to perform"
                },
                "key": {
                    "type": "string",
                    "description": "The key (required for set, get, delete; max 255 chars)"
                },
                "value": {
                    "type": "string",
                    "description": "The value to store (required for set; can be JSON-encoded)"
                }
            },
            "required": ["operation"],
            "additionalProperties": false
        })
    }

    fn hints(&self) -> ToolHints {
        // Mutates shared session storage on set/delete; serialize storage
        // mutations within a batch to avoid lost updates.
        ToolHints::default()
            .with_idempotent(true)
            .with_concurrency_class("session_storage")
    }

    async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
        ToolExecutionResult::tool_error(
            "kv_store requires context. This tool must be executed with session context.",
        )
    }

    async fn execute_with_context(
        &self,
        arguments: Value,
        context: &ToolContext,
    ) -> ToolExecutionResult {
        let operation = match arguments.get("operation").and_then(|v| v.as_str()) {
            Some(op) => op,
            None => {
                return ToolExecutionResult::tool_error("Missing required parameter: operation");
            }
        };

        let storage_store = match &context.storage_store {
            Some(store) => store,
            None => {
                return ToolExecutionResult::tool_error("Storage not available in this context");
            }
        };

        match operation {
            "set" => {
                let key = match arguments.get("key").and_then(|v| v.as_str()) {
                    Some(k) => k,
                    None => {
                        return ToolExecutionResult::tool_error(
                            "Missing required parameter: key (for set operation)",
                        );
                    }
                };
                let value = match arguments.get("value").and_then(|v| v.as_str()) {
                    Some(v) => v,
                    None => {
                        return ToolExecutionResult::tool_error(
                            "Missing required parameter: value (for set operation)",
                        );
                    }
                };
                if key.len() > 255 {
                    return ToolExecutionResult::tool_error("Key must be 255 characters or less");
                }
                match storage_store
                    .set_value(context.session_id, key, value)
                    .await
                {
                    Ok(()) => ToolExecutionResult::success(json!({
                        "operation": "set",
                        "key": key,
                        "success": true
                    })),
                    Err(e) => ToolExecutionResult::internal_error(e),
                }
            }
            "get" => {
                let key = match arguments.get("key").and_then(|v| v.as_str()) {
                    Some(k) => k,
                    None => {
                        return ToolExecutionResult::tool_error(
                            "Missing required parameter: key (for get operation)",
                        );
                    }
                };
                match storage_store.get_value(context.session_id, key).await {
                    Ok(Some(value)) => ToolExecutionResult::success(json!({
                        "operation": "get",
                        "key": key,
                        "value": value,
                        "found": true
                    })),
                    Ok(None) => ToolExecutionResult::success(json!({
                        "operation": "get",
                        "key": key,
                        "value": null,
                        "found": false
                    })),
                    Err(e) => ToolExecutionResult::internal_error(e),
                }
            }
            "delete" => {
                let key = match arguments.get("key").and_then(|v| v.as_str()) {
                    Some(k) => k,
                    None => {
                        return ToolExecutionResult::tool_error(
                            "Missing required parameter: key (for delete operation)",
                        );
                    }
                };
                match storage_store.delete_value(context.session_id, key).await {
                    Ok(deleted) => ToolExecutionResult::success(json!({
                        "operation": "delete",
                        "key": key,
                        "deleted": deleted
                    })),
                    Err(e) => ToolExecutionResult::internal_error(e),
                }
            }
            "list" => match storage_store.list_keys(context.session_id).await {
                Ok(keys) => {
                    let key_list: Vec<Value> = keys
                        .iter()
                        .map(|k| {
                            json!({
                                "key": k.key,
                                "created_at": k.created_at.to_rfc3339(),
                                "updated_at": k.updated_at.to_rfc3339()
                            })
                        })
                        .collect();
                    ToolExecutionResult::success(json!({
                        "operation": "list",
                        "keys": key_list,
                        "count": key_list.len()
                    }))
                }
                Err(e) => ToolExecutionResult::internal_error(e),
            },
            _ => ToolExecutionResult::tool_error(format!(
                "Invalid operation: {}. Must be one of: set, get, delete, list",
                operation
            )),
        }
    }

    fn requires_context(&self) -> bool {
        true
    }
}

// ============================================================================
// SecretStoreTool - Unified secret storage tool
// ============================================================================

/// Tool for encrypted secret storage operations
pub struct SecretStoreTool;

#[async_trait]
impl Tool for SecretStoreTool {
    fn name(&self) -> &str {
        "secret_store"
    }

    fn display_name(&self) -> Option<&str> {
        Some("Secret Store")
    }

    fn description(&self) -> &str {
        "Encrypted secret storage operations: set, get, delete, or list secrets."
    }

    fn parameters_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "enum": ["set", "get", "delete", "list"],
                    "description": "The operation to perform"
                },
                "name": {
                    "type": "string",
                    "description": "The secret name (required for set, get, delete; max 255 chars)"
                },
                "value": {
                    "type": "string",
                    "description": "The secret value to store (required for set; will be encrypted)"
                }
            },
            "required": ["operation"],
            "additionalProperties": false
        })
    }

    fn hints(&self) -> ToolHints {
        // Shares the session storage backend with kv_store; serialize storage
        // mutations within a batch to avoid lost updates.
        ToolHints::default()
            .with_idempotent(true)
            .with_requires_secrets(true)
            .with_concurrency_class("session_storage")
    }

    async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
        ToolExecutionResult::tool_error(
            "secret_store requires context. This tool must be executed with session context.",
        )
    }

    async fn execute_with_context(
        &self,
        arguments: Value,
        context: &ToolContext,
    ) -> ToolExecutionResult {
        let operation = match arguments.get("operation").and_then(|v| v.as_str()) {
            Some(op) => op,
            None => {
                return ToolExecutionResult::tool_error("Missing required parameter: operation");
            }
        };

        let storage_store = match &context.storage_store {
            Some(store) => store,
            None => {
                return ToolExecutionResult::tool_error("Storage not available in this context");
            }
        };

        match operation {
            "set" => {
                let name = match arguments.get("name").and_then(|v| v.as_str()) {
                    Some(n) => n,
                    None => {
                        return ToolExecutionResult::tool_error(
                            "Missing required parameter: name (for set operation)",
                        );
                    }
                };
                let value = match arguments.get("value").and_then(|v| v.as_str()) {
                    Some(v) => v,
                    None => {
                        return ToolExecutionResult::tool_error(
                            "Missing required parameter: value (for set operation)",
                        );
                    }
                };
                if name.len() > 255 {
                    return ToolExecutionResult::tool_error(
                        "Secret name must be 255 characters or less",
                    );
                }
                if is_internal_secret_name(name) {
                    return ToolExecutionResult::tool_error(
                        "Secret name is reserved for internal system use",
                    );
                }
                match storage_store
                    .set_secret(context.session_id, name, value)
                    .await
                {
                    Ok(()) => ToolExecutionResult::success(json!({
                        "operation": "set",
                        "name": name,
                        "success": true
                    })),
                    Err(e) => {
                        let msg = e.to_string();
                        if msg.contains("Encryption not configured") {
                            ToolExecutionResult::tool_error(
                                "Secret storage not available. Encryption is not configured.",
                            )
                        } else {
                            ToolExecutionResult::internal_error(e)
                        }
                    }
                }
            }
            "get" => {
                let name = match arguments.get("name").and_then(|v| v.as_str()) {
                    Some(n) => n,
                    None => {
                        return ToolExecutionResult::tool_error(
                            "Missing required parameter: name (for get operation)",
                        );
                    }
                };
                if is_internal_secret_name(name) {
                    return ToolExecutionResult::tool_error("Secret not found");
                }
                match storage_store.get_secret(context.session_id, name).await {
                    Ok(Some(value)) => ToolExecutionResult::success(json!({
                        "operation": "get",
                        "name": name,
                        "value": value,
                        "found": true
                    })),
                    Ok(None) => ToolExecutionResult::success(json!({
                        "operation": "get",
                        "name": name,
                        "value": null,
                        "found": false
                    })),
                    Err(e) => {
                        let msg = e.to_string();
                        if msg.contains("Encryption not configured") {
                            ToolExecutionResult::tool_error(
                                "Secret storage not available. Encryption is not configured.",
                            )
                        } else {
                            ToolExecutionResult::internal_error(e)
                        }
                    }
                }
            }
            "delete" => {
                let name = match arguments.get("name").and_then(|v| v.as_str()) {
                    Some(n) => n,
                    None => {
                        return ToolExecutionResult::tool_error(
                            "Missing required parameter: name (for delete operation)",
                        );
                    }
                };
                if is_internal_secret_name(name) {
                    return ToolExecutionResult::tool_error(
                        "Secret name is reserved for internal system use",
                    );
                }
                match storage_store.delete_secret(context.session_id, name).await {
                    Ok(deleted) => ToolExecutionResult::success(json!({
                        "operation": "delete",
                        "name": name,
                        "deleted": deleted
                    })),
                    Err(e) => ToolExecutionResult::internal_error(e),
                }
            }
            "list" => match storage_store.list_secrets(context.session_id).await {
                Ok(secrets) => {
                    let secret_list: Vec<Value> = secrets
                        .iter()
                        .filter(|s| !is_internal_secret_name(&s.name))
                        .map(|s| {
                            json!({
                                "name": s.name,
                                "created_at": s.created_at.to_rfc3339(),
                                "updated_at": s.updated_at.to_rfc3339()
                            })
                        })
                        .collect();
                    ToolExecutionResult::success(json!({
                        "operation": "list",
                        "secrets": secret_list,
                        "count": secret_list.len()
                    }))
                }
                Err(e) => {
                    let msg = e.to_string();
                    if msg.contains("Encryption not configured") {
                        ToolExecutionResult::tool_error(
                            "Secret storage not available. Encryption is not configured.",
                        )
                    } else {
                        ToolExecutionResult::internal_error(e)
                    }
                }
            },
            _ => ToolExecutionResult::tool_error(format!(
                "Invalid operation: {}. Must be one of: set, get, delete, list",
                operation
            )),
        }
    }

    fn requires_context(&self) -> bool {
        true
    }
}

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

    #[test]
    fn test_internal_secret_name_filtering() {
        assert!(is_internal_secret_name("browserless_internal:cookies"));
        assert!(!is_internal_secret_name("api_key"));
    }

    #[test]
    fn test_capability_metadata() {
        let cap = SessionStorageCapability;
        assert_eq!(cap.id(), "session_storage");
        assert_eq!(cap.name(), "Storage");
        assert_eq!(cap.status(), CapabilityStatus::Available);
        assert_eq!(cap.icon(), Some("database"));
        assert_eq!(cap.category(), Some("Storage"));
    }

    #[test]
    fn test_capability_has_two_tools() {
        let cap = SessionStorageCapability;
        let tools = cap.tools();

        assert_eq!(tools.len(), 2);

        let tool_names: Vec<&str> = tools.iter().map(|t| t.name()).collect();
        assert!(tool_names.contains(&"kv_store"));
        assert!(tool_names.contains(&"secret_store"));
    }

    #[test]
    fn test_capability_has_system_prompt() {
        let cap = SessionStorageCapability;
        let prompt = cap.system_prompt_addition().unwrap();
        assert!(prompt.contains("kv_store"));
        assert!(prompt.contains("secret_store"));
        assert!(prompt.contains("encrypted"));
    }

    #[test]
    fn test_tools_require_context() {
        assert!(KvStoreTool.requires_context());
        assert!(SecretStoreTool.requires_context());
    }

    #[tokio::test]
    async fn test_kv_store_without_context() {
        let tool = KvStoreTool;
        let result = tool
            .execute(json!({"operation": "set", "key": "test", "value": "data"}))
            .await;

        if let ToolExecutionResult::ToolError(msg) = result {
            assert!(msg.contains("requires context"));
        } else {
            panic!("Expected tool error");
        }
    }

    #[tokio::test]
    async fn test_kv_store_missing_operation() {
        let tool = KvStoreTool;
        let context = ToolContext::new(SessionId::new());

        let result = tool
            .execute_with_context(json!({"key": "test"}), &context)
            .await;

        if let ToolExecutionResult::ToolError(msg) = result {
            assert!(msg.contains("operation"));
        } else {
            panic!("Expected tool error for missing operation");
        }
    }

    #[tokio::test]
    async fn test_kv_store_no_storage_store() {
        let tool = KvStoreTool;
        let context = ToolContext::new(SessionId::new());

        let result = tool
            .execute_with_context(
                json!({"operation": "set", "key": "test", "value": "data"}),
                &context,
            )
            .await;

        if let ToolExecutionResult::ToolError(msg) = result {
            assert!(msg.contains("not available"));
        } else {
            panic!("Expected tool error for missing storage store");
        }
    }

    #[tokio::test]
    async fn test_secret_store_without_context() {
        let tool = SecretStoreTool;
        let result = tool
            .execute(json!({"operation": "set", "name": "api_key", "value": "secret123"}))
            .await;

        if let ToolExecutionResult::ToolError(msg) = result {
            assert!(msg.contains("requires context"));
        } else {
            panic!("Expected tool error");
        }
    }
}