use super::{Capability, CapabilityLocalization, 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_KV_PREFIXES: &[&str] = &[
super::a2a_delegation::AGENT_RUN_KEY_PREFIX,
crate::ard_attachment::ARD_ATTACHMENT_KV_PREFIX,
crate::ard_attachment::ARD_DISCOVERY_KV_PREFIX,
];
const INTERNAL_SECRET_PREFIXES: &[&str] = &["browserless_internal:"];
pub fn is_internal_session_kv_key(key: &str) -> bool {
INTERNAL_KV_PREFIXES
.iter()
.any(|prefix| key.starts_with(prefix))
}
fn reserved_kv_key_error() -> ToolExecutionResult {
ToolExecutionResult::tool_error("Key is reserved for internal system use")
}
fn is_internal_secret_name(name: &str) -> bool {
INTERNAL_SECRET_PREFIXES
.iter()
.any(|prefix| name.starts_with(prefix))
}
pub const SESSION_STORAGE_CAPABILITY_ID: &str = "session_storage";
pub struct SessionStorageCapability;
impl Capability for SessionStorageCapability {
fn id(&self) -> &str {
SESSION_STORAGE_CAPABILITY_ID
}
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 localizations(&self) -> Vec<CapabilityLocalization> {
vec![CapabilityLocalization::text(
"uk",
"Сховище",
r#"Інструменти для збереження та отримання пар ключ-значення і зашифрованих секретів у межах сесії.
> [!NOTE]
> Дані зберігаються протягом усієї сесії. Секрети шифруються при зберіганні.
> [!TIP]
> Використовуйте сховище ключ-значення для загальних даних. Використовуйте секрети для чутливої інформації, як-от API-ключі чи токени."#,
)]
}
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"]
}
}
pub struct KvStoreTool;
#[async_trait]
impl Tool for KvStoreTool {
fn narrate(
&self,
tool_call: &crate::tool_types::ToolCall,
phase: crate::tool_narration::ToolNarrationPhase,
locale: Option<&str>,
) -> Option<String> {
let fallback = self.display_name().unwrap_or("Key-Value Store");
Some(crate::tool_narration::narrate_secret_store(
&tool_call.arguments,
fallback,
phase,
locale,
))
}
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 {
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");
}
if is_internal_session_kv_key(key) {
return reserved_kv_key_error();
}
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)",
);
}
};
if is_internal_session_kv_key(key) {
return reserved_kv_key_error();
}
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)",
);
}
};
if is_internal_session_kv_key(key) {
return reserved_kv_key_error();
}
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()
.filter(|k| !is_internal_session_kv_key(&k.key))
.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
}
}
pub struct SecretStoreTool;
#[async_trait]
impl Tool for SecretStoreTool {
fn narrate(
&self,
tool_call: &crate::tool_types::ToolCall,
phase: crate::tool_narration::ToolNarrationPhase,
locale: Option<&str>,
) -> Option<String> {
let fallback = self.display_name().unwrap_or("Secret Store");
Some(crate::tool_narration::narrate_secret_store(
&tool_call.arguments,
fallback,
phase,
locale,
))
}
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 {
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::traits::SessionStorageStore;
use crate::typed_id::SessionId;
use crate::{KeyInfo, Result};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
#[derive(Default)]
struct TestStorageStore {
values: Mutex<HashMap<String, String>>,
}
#[async_trait]
impl crate::traits::SessionStorageStore for TestStorageStore {
async fn set_value(&self, _session_id: SessionId, key: &str, value: &str) -> Result<()> {
self.values
.lock()
.unwrap()
.insert(key.to_string(), value.to_string());
Ok(())
}
async fn get_value(&self, _session_id: SessionId, key: &str) -> Result<Option<String>> {
Ok(self.values.lock().unwrap().get(key).cloned())
}
async fn delete_value(&self, _session_id: SessionId, key: &str) -> Result<bool> {
Ok(self.values.lock().unwrap().remove(key).is_some())
}
async fn list_keys(&self, _session_id: SessionId) -> Result<Vec<KeyInfo>> {
let now = chrono::Utc::now();
Ok(self
.values
.lock()
.unwrap()
.keys()
.map(|key| KeyInfo {
key: key.clone(),
created_at: now,
updated_at: now,
})
.collect())
}
async fn set_secret(
&self,
_session_id: SessionId,
_name: &str,
_value: &str,
) -> Result<()> {
Ok(())
}
async fn get_secret(&self, _session_id: SessionId, _name: &str) -> Result<Option<String>> {
Ok(None)
}
async fn delete_secret(&self, _session_id: SessionId, _name: &str) -> Result<bool> {
Ok(false)
}
async fn list_secrets(&self, _session_id: SessionId) -> Result<Vec<crate::SecretInfo>> {
Ok(Vec::new())
}
}
#[test]
fn test_internal_kv_key_filtering() {
assert!(is_internal_session_kv_key("agent_run:abc"));
assert!(!is_internal_session_kv_key("user:agent_run:abc"));
}
#[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_kv_store_rejects_reserved_internal_keys() {
let tool = KvStoreTool;
let session_id = SessionId::new();
let storage = Arc::new(TestStorageStore::default());
storage
.set_value(session_id, "agent_run:trusted", "trusted-record")
.await
.unwrap();
storage
.set_value(session_id, "public", "public-record")
.await
.unwrap();
let context = ToolContext::with_storage_store(session_id, storage.clone());
for arguments in [
json!({"operation": "set", "key": "agent_run:trusted", "value": "forged"}),
json!({"operation": "get", "key": "agent_run:trusted"}),
json!({"operation": "delete", "key": "agent_run:trusted"}),
] {
let result = tool.execute_with_context(arguments, &context).await;
assert!(
matches!(result, ToolExecutionResult::ToolError(ref msg) if msg.contains("reserved")),
"expected reserved-key error, got {result:?}"
);
}
assert_eq!(
storage
.get_value(session_id, "agent_run:trusted")
.await
.unwrap()
.as_deref(),
Some("trusted-record")
);
let result = tool
.execute_with_context(json!({"operation": "list"}), &context)
.await;
let ToolExecutionResult::Success(value) = result else {
panic!("expected successful list");
};
assert_eq!(value["count"], 1);
assert_eq!(value["keys"][0]["key"], "public");
}
#[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");
}
}
}