use std::time::Duration;
use super::builder::AgentBuilder;
use super::types::PermissionError;
use super::Agent;
use crate::permission::{
AuthorizationResponse, Grant, GrantStore, Scope, ToolAuthorizationPolicy, ToolCallAuthorizer,
};
impl Agent {
pub fn authorizer(&self) -> &tokio::sync::RwLock<ToolCallAuthorizer> {
&self.authorizer
}
pub async fn respond_to_authorization(
&self,
proposal_id: &str,
response: AuthorizationResponse,
) -> Result<(), PermissionError> {
let pending = self.pending_authorizations.read().await;
if let Some(tx) = pending.get(proposal_id) {
tx.send(response)
.await
.map_err(|_| PermissionError::ChannelClosed)?;
Ok(())
} else {
Err(PermissionError::RequestNotFound(proposal_id.to_string()))
}
}
pub async fn grant_tool_permission(
&self,
proposal_id: &str,
tool_name: &str,
scope: Scope,
) -> Result<(), PermissionError> {
let grant = Grant::tool(tool_name).with_scope(scope);
self.respond_to_authorization(proposal_id, AuthorizationResponse::Trust { grant })
.await
}
pub async fn grant_params_permission(
&self,
proposal_id: &str,
tool_name: &str,
params_hash: &str,
scope: Scope,
) -> Result<(), PermissionError> {
let grant = Grant::exact(tool_name, params_hash).with_scope(scope);
self.respond_to_authorization(proposal_id, AuthorizationResponse::Trust { grant })
.await
}
pub async fn authorize_once(&self, proposal_id: &str) -> Result<(), PermissionError> {
self.respond_to_authorization(proposal_id, AuthorizationResponse::Once)
.await
}
pub async fn deny_authorization(
&self,
proposal_id: &str,
reason: Option<String>,
) -> Result<(), PermissionError> {
self.respond_to_authorization(proposal_id, AuthorizationResponse::Deny { reason })
.await
}
}
impl AgentBuilder {
pub fn with_grant_store(mut self, store: impl GrantStore + 'static) -> Self {
self.grant_store = Some(Box::new(store));
self
}
pub fn with_authorization_timeout(mut self, timeout: Duration) -> Self {
self.authorization_timeout = timeout;
self
}
pub fn interactive(mut self) -> Self {
self.authorization_policy = ToolAuthorizationPolicy::Interactive;
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_builder_authorization_timeout() {
let timeout = Duration::from_secs(60);
let builder = Agent::builder().with_authorization_timeout(timeout);
assert_eq!(builder.authorization_timeout, timeout);
}
#[test]
fn test_builder_grant_store() {
use crate::permission::MemoryGrantStore;
let builder = Agent::builder().with_grant_store(MemoryGrantStore::new());
assert!(builder.grant_store.is_some());
}
}