Complete production readiness: secrets TTL, Cedar reload, audit export, rate limiting

Author: symbibot

crates/runtime/src/api/composio_executor.rs

@@ -6,6 +6,9 @@
 #[cfg(feature = "composio")]
 use std::sync::Arc;
 
+#[cfg(feature = "composio")]
+use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+
 #[cfg(feature = "composio")]
 use async_trait::async_trait;
 
@@ -22,15 +25,78 @@ use crate::reasoning::inference::ToolDefinition;
 #[cfg(feature = "composio")]
 use crate::reasoning::loop_types::{LoopConfig, Observation, ProposedAction};
 
+/// Simple per-minute rate limiter for MCP tool calls.
+#[cfg(feature = "composio")]
+struct McpRateLimiter {
+    /// Maximum calls allowed per minute (0 = unlimited).
+    max_per_minute: u32,
+    /// Calls made in the current window.
+    calls: AtomicU32,
+    /// Start of the current window (seconds since UNIX epoch).
+    window_start: AtomicU64,
+}
+
+#[cfg(feature = "composio")]
+impl McpRateLimiter {
+    fn new(max_per_minute: Option<u32>) -> Self {
+        let now = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_secs();
+        Self {
+            max_per_minute: max_per_minute.unwrap_or(0),
+            calls: AtomicU32::new(0),
+            window_start: AtomicU64::new(now),
+        }
+    }
+
+    /// Check if a call is allowed. Returns `true` if within limits.
+    fn check(&self) -> bool {
+        if self.max_per_minute == 0 {
+            return true; // unlimited
+        }
+
+        let now = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_secs();
+        let window = self.window_start.load(Ordering::Relaxed);
+
+        // Reset window if more than 60 seconds have passed
+        if now - window >= 60 {
+            self.window_start.store(now, Ordering::Relaxed);
+            self.calls.store(1, Ordering::Relaxed);
+            return true;
+        }
+
+        let current = self.calls.fetch_add(1, Ordering::Relaxed);
+        current < self.max_per_minute
+    }
+}
+
 /// An [`ActionExecutor`] that dispatches tool calls to Composio via JSON-RPC.
 #[cfg(feature = "composio")]
 pub struct ComposioToolExecutor {
     transport: Arc<SseTransport>,
     tool_definitions: Vec<ToolDefinition>,
+    rate_limiter: McpRateLimiter,
 }
 
 #[cfg(feature = "composio")]
 impl ComposioToolExecutor {
+    /// Discover available tools from the Composio MCP endpoint and return a
+    /// new executor ready to dispatch calls.
+    ///
+    /// `max_calls_per_minute` enforces a per-server rate limit (None = unlimited).
+    pub async fn discover_with_rate_limit(
+        transport: Arc<SseTransport>,
+        max_calls_per_minute: Option<u32>,
+    ) -> Result<Self, ComposioError> {
+        let mut executor = Self::discover(transport).await?;
+        executor.rate_limiter = McpRateLimiter::new(max_calls_per_minute);
+        Ok(executor)
+    }
+
     /// Discover available tools from the Composio MCP endpoint and return a
     /// new executor ready to dispatch calls.
     pub async fn discover(transport: Arc<SseTransport>) -> Result<Self, ComposioError> {
@@ -69,6 +135,7 @@ impl ComposioToolExecutor {
         Ok(Self {
             transport,
             tool_definitions,
+            rate_limiter: McpRateLimiter::new(None),
         })
     }
 
@@ -79,6 +146,14 @@ impl ComposioToolExecutor {
 
     /// Call a single tool on the Composio MCP endpoint.
     async fn call_tool(&self, name: &str, arguments: &str) -> Result<String, String> {
+        if !self.rate_limiter.check() {
+            tracing::warn!(tool = name, "MCP rate limit exceeded");
+            return Err(format!(
+                "Rate limit exceeded: max {} calls/min for MCP server",
+                self.rate_limiter.max_per_minute
+            ));
+        }
+
         let args: serde_json::Value =
             serde_json::from_str(arguments).unwrap_or(serde_json::json!({}));
 
@@ -199,6 +274,24 @@ mod tests {
         assert!(defs[0].parameters["properties"]["text"].is_object());
     }
 
+    #[test]
+    fn test_rate_limiter_unlimited() {
+        let limiter = McpRateLimiter::new(None);
+        for _ in 0..1000 {
+            assert!(limiter.check());
+        }
+    }
+
+    #[test]
+    fn test_rate_limiter_enforced() {
+        let limiter = McpRateLimiter::new(Some(5));
+        for _ in 0..5 {
+            assert!(limiter.check());
+        }
+        // 6th call should be rejected
+        assert!(!limiter.check());
+    }
+
     #[test]
     fn test_mcp_content_extraction() {
         let result = serde_json::json!({

crates/runtime/src/reasoning/cedar_gate.rs

@@ -89,6 +89,50 @@ impl CedarPolicyGate {
         self.policies.read().await.clone()
     }
 
+    /// Reload policies from a JSON file (hot-reload for production).
+    ///
+    /// The file must contain a JSON array of `CedarPolicy` objects.
+    /// All existing policies are replaced atomically.
+    pub async fn reload_policies_from_file(
+        &self,
+        path: &std::path::Path,
+    ) -> Result<usize, CedarGateError> {
+        let contents = tokio::fs::read_to_string(path).await.map_err(|e| {
+            CedarGateError::ParseError(format!("Failed to read {}: {}", path.display(), e))
+        })?;
+        let new_policies: Vec<CedarPolicy> = serde_json::from_str(&contents).map_err(|e| {
+            CedarGateError::ParseError(format!("Invalid JSON in {}: {}", path.display(), e))
+        })?;
+
+        // Validate all policies parse as valid Cedar before swapping
+        for policy in &new_policies {
+            if policy.active {
+                if let Err(e) = policy.source.parse::<PolicySet>() {
+                    return Err(CedarGateError::ParseError(format!(
+                        "Policy '{}' has invalid Cedar syntax: {}",
+                        policy.name, e
+                    )));
+                }
+            }
+        }
+
+        let count = new_policies.len();
+        let mut policies = self.policies.write().await;
+        *policies = new_policies;
+        tracing::info!(
+            count,
+            path = %path.display(),
+            "Cedar policies reloaded"
+        );
+        Ok(count)
+    }
+
+    /// Replace all policies atomically (for programmatic hot-reload).
+    pub async fn replace_policies(&self, new_policies: Vec<CedarPolicy>) {
+        let mut policies = self.policies.write().await;
+        *policies = new_policies;
+    }
+
     /// Get active policy count.
     pub async fn active_policy_count(&self) -> usize {
         self.policies

crates/runtime/src/reasoning/journal.rs

@@ -166,6 +166,34 @@ impl JournalWriter for DurableJournal {
     }
 }
 
+/// Export all journal entries for an agent as a JSON string for backup.
+pub async fn export_entries(
+    storage: &dyn JournalStorage,
+    agent_id: &AgentId,
+) -> Result<String, JournalError> {
+    let entries = storage.read_entries(agent_id).await?;
+    serde_json::to_string_pretty(&entries)
+        .map_err(|e| JournalError::WriteFailed(format!("Failed to serialize journal entries: {e}")))
+}
+
+/// Import journal entries from a JSON string (restore from backup).
+///
+/// Entries are appended to storage. Callers should compact first if
+/// a clean restore is desired.
+pub async fn import_entries(
+    storage: &dyn JournalStorage,
+    json: &str,
+) -> Result<usize, JournalError> {
+    let entries: Vec<JournalEntry> = serde_json::from_str(json).map_err(|e| {
+        JournalError::ReadFailed(format!("Failed to deserialize journal entries: {e}"))
+    })?;
+    let count = entries.len();
+    for entry in &entries {
+        storage.store(entry).await?;
+    }
+    Ok(count)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -350,4 +378,46 @@ mod tests {
 
         assert_eq!(journal.last_completed_iteration().await.unwrap(), 7);
     }
+
+    #[tokio::test]
+    async fn test_export_entries() {
+        let storage = MemoryJournalStorage::new();
+        let agent = AgentId::new();
+
+        storage.store(&make_entry(agent, 0, 0)).await.unwrap();
+        storage.store(&make_entry(agent, 1, 1)).await.unwrap();
+
+        let json = export_entries(&storage, &agent).await.unwrap();
+        assert!(json.contains("sequence"));
+
+        // Should be valid JSON array
+        let parsed: Vec<JournalEntry> = serde_json::from_str(&json).unwrap();
+        assert_eq!(parsed.len(), 2);
+    }
+
+    #[tokio::test]
+    async fn test_import_entries() {
+        let storage = MemoryJournalStorage::new();
+        let agent = AgentId::new();
+
+        // Create entries, export, then import into a fresh storage
+        storage.store(&make_entry(agent, 0, 0)).await.unwrap();
+        storage.store(&make_entry(agent, 1, 1)).await.unwrap();
+
+        let json = export_entries(&storage, &agent).await.unwrap();
+
+        let fresh_storage = MemoryJournalStorage::new();
+        let count = import_entries(&fresh_storage, &json).await.unwrap();
+        assert_eq!(count, 2);
+
+        let entries = fresh_storage.read_entries(&agent).await.unwrap();
+        assert_eq!(entries.len(), 2);
+    }
+
+    #[tokio::test]
+    async fn test_import_invalid_json() {
+        let storage = MemoryJournalStorage::new();
+        let result = import_entries(&storage, "not valid json").await;
+        assert!(result.is_err());
+    }
 }

crates/runtime/src/reasoning/tracing_spans.rs

@@ -15,6 +15,8 @@ pub struct LoopTracer {
     loop_id: String,
     start: Instant,
     iteration: u32,
+    /// Optional external request ID for correlation with API requests.
+    request_id: Option<String>,
 }
 
 impl LoopTracer {
@@ -41,9 +43,27 @@ impl LoopTracer {
             loop_id,
             start: Instant::now(),
             iteration: 0,
+            request_id: None,
         }
     }
 
+    /// Set an external request ID for correlation with API requests.
+    pub fn with_request_id(mut self, request_id: String) -> Self {
+        tracing::info!(
+            agent_id = %self.agent_id,
+            loop_id = %self.loop_id,
+            request_id = %request_id,
+            "Correlated with API request"
+        );
+        self.request_id = Some(request_id);
+        self
+    }
+
+    /// Get the request ID, if set.
+    pub fn request_id(&self) -> Option<&str> {
+        self.request_id.as_deref()
+    }
+
     /// Begin a new iteration.
     pub fn begin_iteration(&mut self) {
         self.iteration += 1;
@@ -294,6 +314,20 @@ mod tests {
         assert!(parsed.is_some());
     }
 
+    #[test]
+    fn test_with_request_id() {
+        let agent_id = AgentId::new();
+        let tracer = LoopTracer::start(agent_id).with_request_id("req-abc-123".to_string());
+        assert_eq!(tracer.request_id(), Some("req-abc-123"));
+    }
+
+    #[test]
+    fn test_request_id_default_none() {
+        let agent_id = AgentId::new();
+        let tracer = LoopTracer::start(agent_id);
+        assert!(tracer.request_id().is_none());
+    }
+
     #[test]
     fn test_elapsed_increases() {
         let agent_id = AgentId::new();