cots 0.2.0

Cots.ai SDK for Rust. cots::agents::AgentClient — configure once with tenant_id/agent_id, then guard() every governed action through the PEP/PDP before executing it.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

use std::future::Future;
use std::time::{Duration, Instant};

use super::control_plane::ControlPlaneClient;
use super::types::{
    AgentConfig, ApprovalOutcome, GovernedOutcome, GovernedResult, InterceptResult,
    InterceptWireBody, NormalizedAction,
};
use crate::SdkError;

/// The Rust agent client. Configure once with `AgentConfig`, then call
/// `guard()` around every governed action.
#[derive(Debug, Clone)]
pub struct AgentClient {
    config: AgentConfig,
    pub control_plane: ControlPlaneClient,
    http: reqwest::Client,
}

impl AgentClient {
    pub fn new(config: AgentConfig) -> Self {
        let control_plane = ControlPlaneClient::with_api_key(config.control_plane_url.clone(), config.api_key.clone());
        Self { config, control_plane, http: reqwest::Client::new() }
    }

    /// Raw interceptor call. Prefer `guard()` unless you need to handle the decision yourself.
    pub async fn intercept(&self, action: NormalizedAction) -> Result<InterceptResult, SdkError> {
        let body = InterceptWireBody {
            tenant_id: self.config.tenant_id.clone(),
            agent_id: self.config.agent_id.clone(),
            target_system: action.target_system,
            action_type: action.action_type,
            action_name: action.action_name,
            principal_id: action.principal_id,
            session_id: action.session_id,
            data_classification: action.data_classification,
            risk_score: action.risk_score,
            amount: action.amount,
            recipient: action.recipient,
        };
        let url = format!("{}/v1/intercept", self.config.data_plane_url);
        let mut req = self.http.post(&url).json(&body);
        if let Some(key) = &self.config.api_key {
            req = req.header("x-cots-api-key", key);
        }
        let res = req.send().await?;
        let status = res.status();
        let text = res.text().await?;
        if !status.is_success() {
            return Err(SdkError::Api { status: status.as_u16(), body: text });
        }
        serde_json::from_str(&text).map_err(|e| SdkError::Api { status: status.as_u16(), body: e.to_string() })
    }

    /// Poll until a held action_event resolves, or `Timeout`.
    pub async fn wait_for_approval(&self, action_event_id: &str) -> Result<ApprovalOutcome, SdkError> {
        let deadline = Instant::now() + Duration::from_millis(self.config.approval_timeout_ms);
        while Instant::now() < deadline {
            if let Ok(action) = self.control_plane.get_action(action_event_id).await {
                match action.status.as_str() {
                    "executed" => return Ok(ApprovalOutcome::Approved),
                    "denied" | "blocked" => return Ok(ApprovalOutcome::Denied),
                    _ => {}
                }
            }
            tokio::time::sleep(Duration::from_millis(self.config.approval_poll_ms)).await;
        }
        Ok(ApprovalOutcome::Timeout)
    }

    /// Calls the interceptor, then obeys the decision. If allowed, runs
    /// `execute` immediately. If blocked, `execute` never runs. If it
    /// requires approval, holds until a human (or the Approval Center)
    /// decides, then runs `execute` only if approved.
    pub async fn guard<F, Fut, T>(
        &self,
        action: NormalizedAction,
        execute: F,
    ) -> Result<GovernedResult<T>, SdkError>
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = T>,
    {
        let result = self.intercept(action).await?;

        if result.decision == "blocked" {
            return Ok(GovernedResult {
                action_event_id: result.action_event_id,
                decision: result.decision,
                outcome: GovernedOutcome::Blocked,
                matched_rules: result.matched_rules,
                detail: None,
            });
        }
        if result.decision == "allowed" {
            let detail = execute().await;
            return Ok(GovernedResult {
                action_event_id: result.action_event_id,
                decision: result.decision,
                outcome: GovernedOutcome::Executed,
                matched_rules: result.matched_rules,
                detail: Some(detail),
            });
        }

        match self.wait_for_approval(&result.action_event_id).await? {
            ApprovalOutcome::Approved => {
                let detail = execute().await;
                Ok(GovernedResult {
                    action_event_id: result.action_event_id,
                    decision: result.decision,
                    outcome: GovernedOutcome::Executed,
                    matched_rules: result.matched_rules,
                    detail: Some(detail),
                })
            }
            ApprovalOutcome::Denied => Ok(GovernedResult {
                action_event_id: result.action_event_id,
                decision: result.decision,
                outcome: GovernedOutcome::Blocked,
                matched_rules: result.matched_rules,
                detail: None,
            }),
            ApprovalOutcome::Timeout => Ok(GovernedResult {
                action_event_id: result.action_event_id,
                decision: result.decision,
                outcome: GovernedOutcome::Timeout,
                matched_rules: result.matched_rules,
                detail: None,
            }),
        }
    }
}