Skip to main content

agentshield/ir/
tool_surface.rs

1use std::collections::BTreeSet;
2
3use serde::{Deserialize, Serialize};
4
5use super::SourceLocation;
6
7/// A declared tool/function exposed by the extension.
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct ToolSurface {
10    pub name: String,
11    pub description: Option<String>,
12    /// JSON Schema of the tool's input parameters.
13    pub input_schema: Option<serde_json::Value>,
14    /// JSON Schema of the tool's output.
15    pub output_schema: Option<serde_json::Value>,
16    /// Permissions declared by the tool (if any).
17    pub declared_permissions: Vec<DeclaredPermission>,
18    /// Source location where the tool is defined.
19    pub defined_at: Option<SourceLocation>,
20    #[serde(default, skip_serializing_if = "BTreeSet::is_empty")]
21    pub declared_capabilities: BTreeSet<Capability>,
22    #[serde(default, skip_serializing_if = "Vec::is_empty")]
23    pub capability_declarations: Vec<CapabilityDeclaration>,
24    #[serde(default, skip_serializing_if = "BTreeSet::is_empty")]
25    pub observed_capabilities: BTreeSet<Capability>,
26    #[serde(default, skip_serializing_if = "is_false")]
27    pub capability_observation_complete: bool,
28    #[serde(default, skip_serializing_if = "Vec::is_empty")]
29    pub capability_evidence: Vec<CapabilityEvidence>,
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
33#[serde(rename_all = "snake_case")]
34pub enum Capability {
35    FsRead,
36    FsWrite,
37    NetworkEgress,
38    ProcessExec,
39    EnvRead,
40    CredentialAccess,
41    DynamicEval,
42    PackageInstall,
43    DatabaseRead,
44    DatabaseWrite,
45}
46
47impl Capability {
48    pub(crate) fn code(self) -> &'static str {
49        match self {
50            Self::FsRead => "fs_read",
51            Self::FsWrite => "fs_write",
52            Self::NetworkEgress => "network_egress",
53            Self::ProcessExec => "process_exec",
54            Self::EnvRead => "env_read",
55            Self::CredentialAccess => "credential_access",
56            Self::DynamicEval => "dynamic_eval",
57            Self::PackageInstall => "package_install",
58            Self::DatabaseRead => "database_read",
59            Self::DatabaseWrite => "database_write",
60        }
61    }
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65pub struct CapabilityEvidence {
66    pub capability: Capability,
67    pub location: SourceLocation,
68    pub description: String,
69}
70
71#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
72pub struct CapabilityDeclaration {
73    pub capability: Capability,
74    pub source: CapabilityDeclarationSource,
75    pub phrase_or_field: String,
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
79#[serde(rename_all = "snake_case")]
80pub enum CapabilityDeclarationSource {
81    Description,
82    InputSchema,
83    Permission,
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize)]
87pub struct DeclaredPermission {
88    pub permission_type: PermissionType,
89    /// e.g., "filesystem:/tmp/*"
90    pub target: Option<String>,
91    pub description: Option<String>,
92}
93
94#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
95#[serde(rename_all = "snake_case")]
96pub enum PermissionType {
97    FileRead,
98    FileWrite,
99    NetworkAccess,
100    ProcessExec,
101    EnvAccess,
102    DatabaseAccess,
103    Unknown,
104}
105
106fn is_false(value: &bool) -> bool {
107    !*value
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    fn tool() -> ToolSurface {
115        ToolSurface {
116            name: "read_file".into(),
117            description: Some("Read a file".into()),
118            input_schema: None,
119            output_schema: None,
120            declared_permissions: Vec::new(),
121            defined_at: None,
122            declared_capabilities: BTreeSet::new(),
123            capability_declarations: Vec::new(),
124            observed_capabilities: BTreeSet::new(),
125            capability_observation_complete: false,
126            capability_evidence: Vec::new(),
127        }
128    }
129
130    #[test]
131    fn empty_capability_fields_are_omitted_from_json() {
132        let value = serde_json::to_value(tool()).unwrap();
133
134        assert!(value.get("declared_capabilities").is_none());
135        assert!(value.get("capability_declarations").is_none());
136        assert!(value.get("observed_capabilities").is_none());
137        assert!(value.get("capability_observation_complete").is_none());
138        assert!(value.get("capability_evidence").is_none());
139    }
140
141    #[test]
142    fn legacy_tool_json_deserializes_with_capability_defaults() {
143        let value = serde_json::json!({
144            "name": "read_file",
145            "description": "Read a file",
146            "input_schema": null,
147            "output_schema": null,
148            "declared_permissions": [],
149            "defined_at": null
150        });
151
152        let tool: ToolSurface = serde_json::from_value(value).unwrap();
153
154        assert!(tool.declared_capabilities.is_empty());
155        assert!(tool.capability_declarations.is_empty());
156        assert!(tool.observed_capabilities.is_empty());
157        assert!(!tool.capability_observation_complete);
158        assert!(tool.capability_evidence.is_empty());
159    }
160
161    #[test]
162    fn capabilities_serialize_in_stable_enum_order() {
163        let mut tool = tool();
164        tool.observed_capabilities.extend([
165            Capability::NetworkEgress,
166            Capability::FsRead,
167            Capability::CredentialAccess,
168        ]);
169
170        let value = serde_json::to_value(tool).unwrap();
171
172        assert_eq!(
173            value["observed_capabilities"],
174            serde_json::json!(["fs_read", "network_egress", "credential_access"])
175        );
176    }
177
178    #[test]
179    fn populated_capability_state_round_trips() {
180        let mut tool = tool();
181        tool.declared_capabilities.insert(Capability::FsRead);
182        tool.capability_declarations.push(CapabilityDeclaration {
183            capability: Capability::FsRead,
184            source: CapabilityDeclarationSource::Description,
185            phrase_or_field: "read files".into(),
186        });
187        tool.observed_capabilities.insert(Capability::FsRead);
188        tool.capability_observation_complete = true;
189        tool.capability_evidence.push(CapabilityEvidence {
190            capability: Capability::FsRead,
191            location: SourceLocation {
192                file: "src/server.ts".into(),
193                line: 4,
194                column: 2,
195                end_line: Some(4),
196                end_column: Some(12),
197            },
198            description: "file read".into(),
199        });
200
201        let value = serde_json::to_value(&tool).unwrap();
202        let decoded: ToolSurface = serde_json::from_value(value).unwrap();
203
204        assert_eq!(decoded.declared_capabilities, tool.declared_capabilities);
205        assert_eq!(
206            decoded.capability_declarations,
207            tool.capability_declarations
208        );
209        assert_eq!(decoded.observed_capabilities, tool.observed_capabilities);
210        assert!(decoded.capability_observation_complete);
211        assert_eq!(decoded.capability_evidence, tool.capability_evidence);
212    }
213}