feat: new Message type that supports encryption#5
Conversation
Signed-off-by: Aminu Oluwaseun Joshua <seun.aminujoshua@gmail.com>
There was a problem hiding this comment.
Pull request overview
Adds support for sending prompts as JWT-wrapped messages, introducing a new Message variant and JWT decoding logic to transform incoming RPC messages into provider-ready prompts.
Changes:
- Extend
Messagewith anEncryptedvariant and convertMessage -> StringviaTryFrom(including JWT decoding). - Add JWT decoding module, new error variants/status mapping, and unit tests covering JWT decode + message conversion.
- Update agent RPC handler to use the new fallible prompt conversion and adjust error-string assertions.
Reviewed changes
Copilot reviewed 9 out of 10 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| src/message.rs | Adds Message::Encrypted and TryFrom<Message> for String to decode JWT-based prompts |
| src/jwt.rs | Introduces JWT claims struct and decode_jwt helper using HS256 validation |
| src/error.rs | Adds JWT-related error variants and centralizes HTTP status mapping for ApiError |
| src/agent.rs | Switches from Message stringification to fallible prompt extraction (try_into) |
| src/lib.rs | Documents feature/usage and adds internal jwt module |
| src/tests/jwt.rs | Adds unit tests for JWT decoding and Message::Encrypted behavior |
| src/tests/mod.rs | Registers the new JWT test module |
| src/tests/unit_error.rs | Updates expected ApiError string formatting |
| Cargo.toml | Adds jsonwebtoken dependency and serial_test dev-dependency |
| Cargo.lock | Locks new dependency graph for JWT + serial test |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Signed-off-by: Aminu Oluwaseun Joshua <seun.aminujoshua@gmail.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 10 changed files in this pull request and generated 5 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| impl TryFrom<Message> for String { | ||
| type Error = ApiError; | ||
|
|
||
| fn try_from(message: Message) -> Result<Self, Self::Error> { | ||
| match message { | ||
| Message::Text(text) => Ok(text), | ||
| Message::Struct(value) => { |
There was a problem hiding this comment.
This PR removes Display formatting for Message and replaces it with TryFrom<Message> for String, which is a breaking change for downstream users relying on message.to_string()/formatting. Consider keeping a Display impl for non-fallible variants (or a non-breaking as_prompt()/try_into_prompt() API) and limiting fallible conversion to an explicit method/trait for the encrypted/JWT case.
| Message::Encrypted(token) => { | ||
| let hmac_secret = | ||
| std::env::var("JWT_SECRET").map_err(|_| Error::NoJWTSecretFound)?; | ||
|
|
||
| let prompt = decode_jwt(&token, &hmac_secret).map(|c| c.prompt)?; | ||
|
|
There was a problem hiding this comment.
Reading JWT_SECRET directly from the process environment here introduces a hidden global dependency and makes prompt extraction harder to configure/test (and requires serializing tests that touch env). Prefer passing the secret via configuration (e.g., AgentServerBuilder) or including it in the request context so conversion doesn’t depend on global env state.
| /// Guard to ensure JWT_SECRET is removed from the environment after test, even on panic. | ||
| struct JwtSecretGuard; | ||
|
|
||
| impl JwtSecretGuard { | ||
| fn set(secret: &str) -> Self { | ||
| unsafe { env::set_var("JWT_SECRET", secret) }; | ||
| JwtSecretGuard | ||
| } | ||
| } | ||
|
|
||
| impl Drop for JwtSecretGuard { | ||
| fn drop(&mut self) { | ||
| unsafe { env::remove_var("JWT_SECRET") }; |
There was a problem hiding this comment.
JwtSecretGuard always removes JWT_SECRET on drop rather than restoring any pre-existing value. If JWT_SECRET is set in the developer/CI environment, these tests will permanently unset it for the remainder of the test process. Consider capturing the previous value in the guard and restoring it in Drop.
| /// Guard to ensure JWT_SECRET is removed from the environment after test, even on panic. | |
| struct JwtSecretGuard; | |
| impl JwtSecretGuard { | |
| fn set(secret: &str) -> Self { | |
| unsafe { env::set_var("JWT_SECRET", secret) }; | |
| JwtSecretGuard | |
| } | |
| } | |
| impl Drop for JwtSecretGuard { | |
| fn drop(&mut self) { | |
| unsafe { env::remove_var("JWT_SECRET") }; | |
| /// Guard to ensure JWT_SECRET is restored after test, even on panic. | |
| struct JwtSecretGuard { | |
| previous: Option<String>, | |
| } | |
| impl JwtSecretGuard { | |
| fn set(secret: &str) -> Self { | |
| let previous = env::var("JWT_SECRET").ok(); | |
| unsafe { env::set_var("JWT_SECRET", secret) }; | |
| JwtSecretGuard { previous } | |
| } | |
| } | |
| impl Drop for JwtSecretGuard { | |
| fn drop(&mut self) { | |
| match &self.previous { | |
| Some(previous) => unsafe { env::set_var("JWT_SECRET", previous) }, | |
| None => unsafe { env::remove_var("JWT_SECRET") }, | |
| } |
| #[test] | ||
| #[serial_test::serial] | ||
| fn message_encrypted_variant_no_secret() { | ||
| unsafe { env::remove_var("JWT_SECRET") }; | ||
| let msg = Message::Encrypted("sometoken".to_string()); | ||
| let result: Result<String, ApiError> = msg.try_into(); | ||
| assert_eq!( | ||
| result.unwrap_err().to_string(), | ||
| ApiError::from(Error::NoJWTSecretFound).to_string() | ||
| ); |
There was a problem hiding this comment.
This test mutates global process environment (remove_var("JWT_SECRET")) without restoring the previous value afterward. Even with serial_test, this can make subsequent tests depend on ordering and can break local runs where JWT_SECRET is set. Prefer using a guard that restores the prior value at the end of the test.
Signed-off-by: Aminu Oluwaseun Joshua <seun.aminujoshua@gmail.com>
Fixes #4