From 144aa7c12c68cd57cd55278736ff87a39dfe374a Mon Sep 17 00:00:00 2001 From: Aminu Oluwaseun Joshua Date: Tue, 7 Apr 2026 13:27:04 +0100 Subject: [PATCH 1/3] makes tracing easier Signed-off-by: Aminu Oluwaseun Joshua --- Cargo.lock | 1 + Cargo.toml | 2 ++ src/agent.rs | 3 ++- src/jwt.rs | 1 + src/providers/ollama.rs | 1 + src/providers/openai.rs | 1 + 6 files changed, 8 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 34d6d22..d878224 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1379,6 +1379,7 @@ dependencies = [ "tarpc", "thiserror 2.0.18", "tokio", + "tracing", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 4b0121b..98a248c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,8 @@ schemars = "1.2.1" thiserror = "2.0.18" +tracing = "0.1.44" + jsonwebtoken = "9.3.1" [dev-dependencies] diff --git a/src/agent.rs b/src/agent.rs index 03a2bfd..9ce8c30 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -28,7 +28,7 @@ impl AgentServer { tarpc::serde_transport::tcp::listen(self.socket_addr, Json::default).await?; listener.config_mut().max_frame_length(usize::MAX); - println!("Listening on: {}", listener.local_addr()); + tracing::info!("Listening on: {}", listener.local_addr()); listener .map_err(|e| eprintln!("{}", e)) // TODO: Improve error handling. @@ -73,6 +73,7 @@ pub(crate) trait AgentWorker { impl AgentWorker for AgentServer { /// Handles a user message by passing it to the completion provider and returning the response. + #[tracing::instrument(name = "agent.message", skip(self, _context, user_message))] async fn message( self, _context: ::tarpc::context::Context, diff --git a/src/jwt.rs b/src/jwt.rs index 5ea8824..d7640a7 100644 --- a/src/jwt.rs +++ b/src/jwt.rs @@ -5,6 +5,7 @@ use crate::error::Error; /// JWT token decoder /// Decodes a JWT token and returns the claims if valid. +#[tracing::instrument(name = "jwt.decode", skip(token, hmac_secret))] pub(crate) fn decode_jwt(token: &str, hmac_secret: &str) -> Result { let validation = jsonwebtoken::Validation::new(Algorithm::HS256); diff --git a/src/providers/ollama.rs b/src/providers/ollama.rs index d58acca..c807157 100644 --- a/src/providers/ollama.rs +++ b/src/providers/ollama.rs @@ -89,6 +89,7 @@ impl OllamaAI { #[async_trait::async_trait] impl CompletionProvider for OllamaAI { + #[tracing::instrument(name = "ollama.chat", skip(self, prompt))] async fn chat(&self, prompt: &str) -> Result { let response = self.agent.prompt(prompt).await?; Ok(response) diff --git a/src/providers/openai.rs b/src/providers/openai.rs index 71e17fe..c6a782c 100644 --- a/src/providers/openai.rs +++ b/src/providers/openai.rs @@ -104,6 +104,7 @@ impl OpenAI { #[async_trait::async_trait] impl CompletionProvider for OpenAI { + #[tracing::instrument(name = "openai.chat", skip(self, prompt))] async fn chat(&self, prompt: &str) -> Result { let response = self.agent.prompt(prompt).await?; Ok(response) From 2d4bd67e40121cfbd11a4b0f6005c47c93bdc3df Mon Sep 17 00:00:00 2001 From: Aminu Oluwaseun Joshua Date: Tue, 7 Apr 2026 13:54:46 +0100 Subject: [PATCH 2/3] make an optional feature Signed-off-by: Aminu Oluwaseun Joshua --- Cargo.toml | 5 ++++- README.md | 4 ++++ src/agent.rs | 6 +++++- src/jwt.rs | 5 ++++- src/lib.rs | 4 ++++ src/providers/ollama.rs | 5 ++++- src/providers/openai.rs | 5 ++++- 7 files changed, 29 insertions(+), 5 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 98a248c..1cf550f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,9 +23,12 @@ schemars = "1.2.1" thiserror = "2.0.18" -tracing = "0.1.44" +tracing = { version = "0.1.44", optional = true } jsonwebtoken = "9.3.1" +[features] +tracing = ["dep:tracing"] + [dev-dependencies] serial_test = "3.4.0" diff --git a/README.md b/README.md index 3b491ec..304cea3 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,10 @@ async fn main() -> Result<(), Box> { --- +## Feature Flag + +**`tracing`**: Enables tracing for logging request/response traces. It requires a `RUST_LOG` environment variable to be set. + ## Contributing Pull requests and issues are welcome! Please open an issue to discuss your ideas or report bugs. diff --git a/src/agent.rs b/src/agent.rs index 9ce8c30..8c3281e 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -28,6 +28,7 @@ impl AgentServer { tarpc::serde_transport::tcp::listen(self.socket_addr, Json::default).await?; listener.config_mut().max_frame_length(usize::MAX); + #[cfg(feature = "tracing")] tracing::info!("Listening on: {}", listener.local_addr()); listener @@ -73,7 +74,10 @@ pub(crate) trait AgentWorker { impl AgentWorker for AgentServer { /// Handles a user message by passing it to the completion provider and returning the response. - #[tracing::instrument(name = "agent.message", skip(self, _context, user_message))] + #[cfg_attr( + feature = "tracing", + tracing::instrument(name = "agent.message", skip(self, _context, user_message)) + )] async fn message( self, _context: ::tarpc::context::Context, diff --git a/src/jwt.rs b/src/jwt.rs index d7640a7..2b8be14 100644 --- a/src/jwt.rs +++ b/src/jwt.rs @@ -5,7 +5,10 @@ use crate::error::Error; /// JWT token decoder /// Decodes a JWT token and returns the claims if valid. -#[tracing::instrument(name = "jwt.decode", skip(token, hmac_secret))] +#[cfg_attr( + feature = "tracing", + tracing::instrument(name = "jwt.decode", skip(token, hmac_secret)) +)] pub(crate) fn decode_jwt(token: &str, hmac_secret: &str) -> Result { let validation = jsonwebtoken::Validation::new(Algorithm::HS256); diff --git a/src/lib.rs b/src/lib.rs index 6ec7061..78c11e8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,6 +28,10 @@ //! Ok(()) //!} //!``` +//! +//! ## Feature Flags +//! +//! **`tracing`**: Enables tracing for logging request/response traces. It requires a `RUST_LOG` environment variable to be set. mod agent; mod builder; diff --git a/src/providers/ollama.rs b/src/providers/ollama.rs index c807157..4771050 100644 --- a/src/providers/ollama.rs +++ b/src/providers/ollama.rs @@ -89,7 +89,10 @@ impl OllamaAI { #[async_trait::async_trait] impl CompletionProvider for OllamaAI { - #[tracing::instrument(name = "ollama.chat", skip(self, prompt))] + #[cfg_attr( + feature = "tracing", + tracing::instrument(name = "ollama.chat", skip(self, prompt)) + )] async fn chat(&self, prompt: &str) -> Result { let response = self.agent.prompt(prompt).await?; Ok(response) diff --git a/src/providers/openai.rs b/src/providers/openai.rs index c6a782c..0982fd4 100644 --- a/src/providers/openai.rs +++ b/src/providers/openai.rs @@ -104,7 +104,10 @@ impl OpenAI { #[async_trait::async_trait] impl CompletionProvider for OpenAI { - #[tracing::instrument(name = "openai.chat", skip(self, prompt))] + #[cfg_attr( + feature = "tracing", + tracing::instrument(name = "openai.chat", skip(self, prompt)) + )] async fn chat(&self, prompt: &str) -> Result { let response = self.agent.prompt(prompt).await?; Ok(response) From feefd98222d4d95526684199b213131ae0896f1c Mon Sep 17 00:00:00 2001 From: Aminu Oluwaseun Joshua Date: Tue, 7 Apr 2026 14:07:01 +0100 Subject: [PATCH 3/3] improved docs Signed-off-by: Aminu Oluwaseun Joshua --- src/agent.rs | 3 +++ src/lib.rs | 4 ++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/agent.rs b/src/agent.rs index 8c3281e..520910f 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -31,6 +31,9 @@ impl AgentServer { #[cfg(feature = "tracing")] tracing::info!("Listening on: {}", listener.local_addr()); + #[cfg(not(feature = "tracing"))] + println!("Listening on: {}", listener.local_addr()); + listener .map_err(|e| eprintln!("{}", e)) // TODO: Improve error handling. .filter_map(|r| future::ready(r.ok())) diff --git a/src/lib.rs b/src/lib.rs index 78c11e8..801e640 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -31,8 +31,8 @@ //! //! ## Feature Flags //! -//! **`tracing`**: Enables tracing for logging request/response traces. It requires a `RUST_LOG` environment variable to be set. - +//! `tracing`: Enables `tracing` instrumentation for request/response traces. +//! mod agent; mod builder; pub mod error;