feat: add core telemetry and trace tooling (#14)

* feat: add core telemetry and trace tooling

Signed-off-by: Luca Muscariello <[email protected]>

* test: cover trace helpers

Signed-off-by: Luca Muscariello <[email protected]>

* test: cover telemetry helpers

Signed-off-by: Luca Muscariello <[email protected]>

* test: cover trace list and summary

Signed-off-by: Luca Muscariello <[email protected]>

* test: expand telemetry and trace coverage

Signed-off-by: Luca Muscariello <[email protected]>

* test: serialize telemetry env tests

Signed-off-by: Luca Muscariello <[email protected]>

---------

Signed-off-by: Luca Muscariello <[email protected]>

Author: muscariello

Cargo.lock

@@ -6,6 +6,7 @@ members = [
   "crates/shadi_memory",
   "crates/shadi_py",
   "crates/shadi_sandbox",
+  "crates/shadi_telemetry",
   "crates/shadictl"
 ]
 resolver = "2"

Cargo.toml

@@ -27,6 +27,7 @@
     _resource = Resource(
         attributes={
             RES_SERVICE_NAME: SERVICE_NAME,
+            "service.namespace": "shadi",
             "telemetry.sdk.language": "python",
         }
     )

agents/secops/telemetry.py

@@ -13,6 +13,8 @@ doctest = false
 agent_secrets = { path = "../agent_secrets", features = ["onepassword"] }
 shadi_memory = { path = "../shadi_memory" }
 shadi_sandbox = { path = "../shadi_sandbox" }
+shadi_telemetry = { path = "../shadi_telemetry" }
+tracing = "0.1"
 
 [dependencies.pyo3]
 version = "0.21"

crates/shadi_py/Cargo.toml

@@ -14,6 +14,7 @@ use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyModule};
 use shadi_memory::{MemoryEntry as ShadiMemoryEntry, SqlCipherStore};
 use shadi_sandbox::{spawn_sandboxed, SandboxError, SandboxPolicy};
+use tracing::{field, info_span};
 
 struct SessionFlagVerifier;
 
@@ -127,6 +128,8 @@ impl ShadiStore {
     }
 
     fn put(&self, session: &PySessionContext, key: &str, secret: &[u8]) -> PyResult<()> {
+        let span = info_span!("shadi.secret.put", secret.key = %key);
+        let _guard = span.enter();
         let ctx = session.to_context();
         let guard = self.store.lock().map_err(|_| PyRuntimeError::new_err("lock poisoned"))?;
         let access = AgentSecretAccess::new(guard.as_ref(), &self.verifier);
@@ -141,6 +144,8 @@ impl ShadiStore {
         session: &PySessionContext,
         key: &str,
     ) -> PyResult<Bound<'py, PyBytes>> {
+        let span = info_span!("shadi.secret.get", secret.key = %key);
+        let _guard = span.enter();
         let ctx = session.to_context();
         let guard = self.store.lock().map_err(|_| PyRuntimeError::new_err("lock poisoned"))?;
         let access = AgentSecretAccess::new(guard.as_ref(), &self.verifier);
@@ -150,6 +155,8 @@ impl ShadiStore {
     }
 
     fn delete(&self, session: &PySessionContext, key: &str) -> PyResult<()> {
+        let span = info_span!("shadi.secret.delete", secret.key = %key);
+        let _guard = span.enter();
         let ctx = session.to_context();
         let guard = self.store.lock().map_err(|_| PyRuntimeError::new_err("lock poisoned"))?;
         let access = AgentSecretAccess::new(guard.as_ref(), &self.verifier);
@@ -159,6 +166,8 @@ impl ShadiStore {
     }
 
     fn list_keys(&self, session: &PySessionContext) -> PyResult<Vec<String>> {
+        let span = info_span!("shadi.secret.list_keys");
+        let _guard = span.enter();
         let ctx = session.to_context();
         AgentSecretAccess::require_verified(&ctx).map_err(map_secret_error)?;
         let guard = self.store.lock().map_err(|_| PyRuntimeError::new_err("lock poisoned"))?;
@@ -178,12 +187,16 @@ impl SqlCipherMemoryStore {
     }
 
     fn put(&self, scope: &str, entry_key: &str, payload: &str) -> PyResult<i64> {
+        let span = info_span!("shadi.memory.put", memory.scope = %scope, memory.entry_key = %entry_key);
+        let _guard = span.enter();
         self.store
             .put(scope, entry_key, payload)
             .map_err(|err| PyRuntimeError::new_err(err.to_string()))
     }
 
     fn get_latest(&self, scope: &str, entry_key: &str) -> PyResult<Option<MemoryEntry>> {
+        let span = info_span!("shadi.memory.get_latest", memory.scope = %scope, memory.entry_key = %entry_key);
+        let _guard = span.enter();
         let entry = self
             .store
             .get_latest(scope, entry_key)
@@ -193,6 +206,13 @@ impl SqlCipherMemoryStore {
 
     #[pyo3(signature = (query, scope=None, limit=10))]
     fn search(&self, query: &str, scope: Option<String>, limit: usize) -> PyResult<Vec<MemoryEntry>> {
+        let span = info_span!(
+            "shadi.memory.search",
+            memory.query = %query,
+            memory.scope = %scope.as_deref().unwrap_or(""),
+            memory.limit = limit as i64,
+        );
+        let _guard = span.enter();
         let entries = self
             .store
             .search(scope.as_deref(), query, limit)
@@ -205,6 +225,12 @@ impl SqlCipherMemoryStore {
 
     #[pyo3(signature = (scope=None, limit=50))]
     fn list(&self, scope: Option<String>, limit: usize) -> PyResult<Vec<MemoryEntry>> {
+        let span = info_span!(
+            "shadi.memory.list",
+            memory.scope = %scope.as_deref().unwrap_or(""),
+            memory.limit = limit as i64,
+        );
+        let _guard = span.enter();
         let entries = self
             .store
             .list(scope.as_deref(), limit)
@@ -216,6 +242,8 @@ impl SqlCipherMemoryStore {
     }
 
     fn delete(&self, scope: &str, entry_key: &str) -> PyResult<usize> {
+        let span = info_span!("shadi.memory.delete", memory.scope = %scope, memory.entry_key = %entry_key);
+        let _guard = span.enter();
         self.store
             .delete(scope, entry_key)
             .map_err(|err| PyRuntimeError::new_err(err.to_string()))
@@ -286,6 +314,7 @@ impl PySessionContext {
 
 #[pymodule]
 fn shadi(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    shadi_telemetry::init("shadi-runtime");
     m.add_class::<ShadiStore>()?;
     m.add_class::<PySessionContext>()?;
     m.add_class::<SqlCipherMemoryStore>()?;
@@ -325,6 +354,8 @@ fn inject_keychain_with_store(
     command: &mut Command,
     mappings: &[String],
 ) -> Result<(), String> {
+    let span = info_span!("shadi.secrets.inject", secret.count = mappings.len() as i64);
+    let _guard = span.enter();
     for mapping in mappings {
         let (key, env) = parse_key_env(mapping)?;
         let secret = store
@@ -361,6 +392,15 @@ fn run_sandboxed(
         return Err(PyRuntimeError::new_err("command must not be empty"));
     }
 
+    let cwd_value = cwd.as_deref().unwrap_or("");
+    let span = info_span!(
+        "shadi.sandbox.run",
+        command = %command[0],
+        cwd = %cwd_value,
+        exit.code = field::Empty,
+    );
+    let _guard = span.enter();
+
     let mut cmd = Command::new(&command[0]);
     if command.len() > 1 {
         cmd.args(&command[1..]);
@@ -381,6 +421,7 @@ fn run_sandboxed(
     let status = child
         .wait()
         .map_err(|err| PyRuntimeError::new_err(err.to_string()))?;
+    span.record("exit.code", &status.code().unwrap_or(-1));
     Ok(status.code().unwrap_or(1))
 }
 

crates/shadi_py/src/lib.rs

@@ -9,6 +9,7 @@ coverage = []
 
 [dependencies]
 thiserror = "1.0"
+tracing = "0.1"
 
 [target.'cfg(target_os = "macos")'.dependencies]
 libc = "0.2"

crates/shadi_sandbox/Cargo.toml

@@ -7,8 +7,32 @@ mod platform;
 pub use policy::SandboxPolicy;
 use std::process::{Command, ExitStatus};
 use std::io;
+use tracing::{field, info_span};
 
 pub fn spawn_sandboxed(command: &mut Command, policy: &SandboxPolicy) -> Result<SandboxedChild, SandboxError> {
+    let program = command.get_program().to_string_lossy().to_string();
+    let args = command
+        .get_args()
+        .map(|arg| arg.to_string_lossy())
+        .collect::<Vec<_>>()
+        .join(" ");
+    let cwd = command
+        .get_current_dir()
+        .map(|path| path.display().to_string())
+        .unwrap_or_else(|| "".to_string());
+    let allowed_paths = policy.allow_read().len() + policy.allow_write().len();
+    let network_mode = if policy.net_blocked() { "blocked" } else { "allowed" };
+
+    let span = info_span!(
+        "shadi.sandbox.spawn",
+        command = %program,
+        args = %args,
+        cwd = %cwd,
+        policy.allowed_paths = allowed_paths as i64,
+        network.mode = %network_mode,
+    );
+    let _guard = span.enter();
+
     platform::spawn_sandboxed(command, policy)
 }
 
@@ -37,14 +61,26 @@ impl SandboxedChild {
     }
 
     pub fn wait(&mut self) -> io::Result<ExitStatus> {
-        match &mut self.inner {
+        let span = info_span!("shadi.sandbox.wait", pid = self.id(), exit.code = field::Empty);
+        let _guard = span.enter();
+
+        let status = match &mut self.inner {
             SandboxedChildInner::Std(child) => child.wait(),
             #[cfg(target_os = "windows")]
             SandboxedChildInner::Windows(child) => child.wait(),
+        };
+
+        if let Ok(ref status) = status {
+            span.record("exit.code", &status.code().unwrap_or(-1));
         }
+
+        status
     }
 
     pub fn kill(&mut self) -> io::Result<()> {
+        let span = info_span!("shadi.sandbox.kill", pid = self.id());
+        let _guard = span.enter();
+
         match &mut self.inner {
             SandboxedChildInner::Std(child) => child.kill(),
             #[cfg(target_os = "windows")]

crates/shadi_sandbox/src/lib.rs

@@ -0,0 +1,13 @@
+[package]
+name = "shadi_telemetry"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+opentelemetry = "0.24"
+opentelemetry_sdk = "0.24"
+opentelemetry-otlp = { version = "0.17", features = ["http-proto", "reqwest-client"] }
+tracing = "0.1"
+tracing-opentelemetry = "0.25"
+tracing-subscriber = { version = "0.3", features = ["json"] }
+tracing-appender = "0.2"