1use agent_sdk_core::{
7 AgentError, CapabilityId, CapabilityNamespace, CapabilityPermission, ExecutorRef,
8 PackageSidecarRef, PolicyKind, PolicyRef, PrivacyClass, SourceRef, ToolPackId, ToolPackKind,
9 ToolPackSnapshot, ToolPackToolSnapshot, TrustClass,
10 policy::{EffectClass, RiskClass},
11 tool_records::CanonicalToolName,
12};
13
14use crate::packs::ToolkitPackBundle;
15
16#[derive(Clone, Debug, Eq, PartialEq)]
17pub enum ToolkitToolExecutionMode {
21 Sync,
23 Async,
25}
26
27#[derive(Clone, Debug, Eq, PartialEq)]
28pub struct Tool {
31 snapshot: ToolPackToolSnapshot,
32 mode: ToolkitToolExecutionMode,
33}
34
35impl Tool {
36 pub fn builder(
38 tool_name: impl Into<String>,
39 executor_ref: impl Into<String>,
40 schema_id: impl Into<String>,
41 policy_ref: PolicyRef,
42 ) -> ToolBuilder {
43 ToolBuilder::new(tool_name, executor_ref, schema_id, policy_ref)
44 }
45
46 pub fn canonical_tool_name(&self) -> &CanonicalToolName {
48 &self.snapshot.canonical_tool_name
49 }
50
51 pub fn executor_ref(&self) -> &ExecutorRef {
53 &self.snapshot.executor_ref
54 }
55
56 pub fn mode(&self) -> &ToolkitToolExecutionMode {
58 &self.mode
59 }
60
61 pub fn snapshot(&self) -> &ToolPackToolSnapshot {
63 &self.snapshot
64 }
65
66 pub fn into_snapshot(self) -> ToolPackToolSnapshot {
68 self.snapshot
69 }
70}
71
72#[derive(Clone, Debug, Eq, PartialEq)]
73pub struct AsyncTool {
77 snapshot: ToolPackToolSnapshot,
78 mode: ToolkitToolExecutionMode,
79}
80
81impl AsyncTool {
82 pub fn builder(
84 tool_name: impl Into<String>,
85 executor_ref: impl Into<String>,
86 schema_id: impl Into<String>,
87 policy_ref: PolicyRef,
88 ) -> ToolBuilder {
89 ToolBuilder::new(tool_name, executor_ref, schema_id, policy_ref).async_mode()
90 }
91
92 pub fn canonical_tool_name(&self) -> &CanonicalToolName {
94 &self.snapshot.canonical_tool_name
95 }
96
97 pub fn executor_ref(&self) -> &ExecutorRef {
99 &self.snapshot.executor_ref
100 }
101
102 pub fn mode(&self) -> &ToolkitToolExecutionMode {
104 &self.mode
105 }
106
107 pub fn snapshot(&self) -> &ToolPackToolSnapshot {
109 &self.snapshot
110 }
111
112 pub fn into_snapshot(self) -> ToolPackToolSnapshot {
114 self.snapshot
115 }
116}
117
118#[derive(Clone, Debug)]
119pub struct ToolBuilder {
123 tool_name: String,
124 executor_ref: String,
125 schema_id: String,
126 description: Option<String>,
127 redacted_schema: Option<serde_json::Value>,
128 policy_refs: Vec<PolicyRef>,
129 requires_approval: bool,
130 capability_id: Option<CapabilityId>,
131 namespace: Option<CapabilityNamespace>,
132 required_permissions: Vec<CapabilityPermission>,
133 effect_class: Option<EffectClass>,
134 risk_class: Option<RiskClass>,
135 redaction_policy_ref: PolicyRef,
136 timeout_ms: u64,
137 cancellation: String,
138 reconciliation: String,
139 privacy: PrivacyClass,
140 mode: ToolkitToolExecutionMode,
141}
142
143impl ToolBuilder {
144 pub fn new(
146 tool_name: impl Into<String>,
147 executor_ref: impl Into<String>,
148 schema_id: impl Into<String>,
149 policy_ref: PolicyRef,
150 ) -> Self {
151 Self {
152 tool_name: tool_name.into(),
153 executor_ref: executor_ref.into(),
154 schema_id: schema_id.into(),
155 description: None,
156 redacted_schema: None,
157 policy_refs: vec![policy_ref],
158 requires_approval: false,
159 capability_id: None,
160 namespace: None,
161 required_permissions: Vec::new(),
162 effect_class: None,
163 risk_class: None,
164 redaction_policy_ref: PolicyRef::with_kind(
165 PolicyKind::Redaction,
166 "policy.redaction.tool_result.refs_only",
167 ),
168 timeout_ms: 10_000,
169 cancellation: "best_effort".to_string(),
170 reconciliation: "effect_lineage_required".to_string(),
171 privacy: PrivacyClass::ContentRefsOnly,
172 mode: ToolkitToolExecutionMode::Sync,
173 }
174 }
175
176 pub fn read_only(mut self) -> Self {
178 self.effect_class = Some(EffectClass::Read);
179 self.risk_class = Some(RiskClass::Low);
180 self
181 }
182
183 pub fn write_effect(mut self) -> Self {
185 self.effect_class = Some(EffectClass::Write);
186 self.risk_class = Some(RiskClass::High);
187 self
188 }
189
190 pub fn effect(mut self, effect_class: EffectClass, risk_class: RiskClass) -> Self {
192 self.effect_class = Some(effect_class);
193 self.risk_class = Some(risk_class);
194 self
195 }
196
197 pub fn capability_id(mut self, capability_id: CapabilityId) -> Self {
199 self.capability_id = Some(capability_id);
200 self
201 }
202
203 pub fn namespace(mut self, namespace: CapabilityNamespace) -> Self {
205 self.namespace = Some(namespace);
206 self
207 }
208
209 pub fn required_permission(mut self, permission: CapabilityPermission) -> Self {
211 self.required_permissions.push(permission);
212 self
213 }
214
215 pub fn redacted_schema(mut self, schema: serde_json::Value) -> Self {
217 self.redacted_schema = Some(schema);
218 self
219 }
220
221 pub fn description(mut self, description: impl Into<String>) -> Self {
223 let description = description.into();
224 if !description.trim().is_empty() {
225 self.description = Some(description);
226 }
227 self
228 }
229
230 pub fn description_opt(mut self, description: Option<String>) -> Self {
232 self.description = description.filter(|description| !description.trim().is_empty());
233 self
234 }
235
236 pub fn policy_ref(mut self, policy_ref: PolicyRef) -> Self {
238 self.policy_refs.push(policy_ref);
239 self
240 }
241
242 pub fn require_approval(mut self) -> Self {
244 self.requires_approval = true;
245 self.policy_refs.push(PolicyRef::with_kind(
246 PolicyKind::Approval,
247 format!("policy.approval.{}", self.tool_name),
248 ));
249 self.risk_class = Some(RiskClass::High);
250 self
251 }
252
253 pub fn redaction_policy(mut self, policy_ref: PolicyRef) -> Self {
255 self.redaction_policy_ref = policy_ref;
256 self
257 }
258
259 pub fn timeout_ms(mut self, timeout_ms: u64) -> Self {
261 self.timeout_ms = timeout_ms;
262 self
263 }
264
265 pub fn cancellation(mut self, cancellation: impl Into<String>) -> Self {
267 self.cancellation = cancellation.into();
268 self
269 }
270
271 pub fn reconciliation(mut self, reconciliation: impl Into<String>) -> Self {
273 self.reconciliation = reconciliation.into();
274 self
275 }
276
277 pub fn privacy(mut self, privacy: PrivacyClass) -> Self {
279 self.privacy = privacy;
280 self
281 }
282
283 pub fn async_mode(mut self) -> Self {
285 self.mode = ToolkitToolExecutionMode::Async;
286 self.timeout_ms = self.timeout_ms.max(60_000);
287 self.cancellation = "cooperative".to_string();
288 self
289 }
290
291 pub fn build(self) -> Result<Tool, AgentError> {
293 let mode = self.mode.clone();
294 Ok(Tool {
295 snapshot: self.build_snapshot()?,
296 mode,
297 })
298 }
299
300 pub fn build_async(self) -> Result<AsyncTool, AgentError> {
302 let builder = self.async_mode();
303 let mode = builder.mode.clone();
304 Ok(AsyncTool {
305 snapshot: builder.build_snapshot()?,
306 mode,
307 })
308 }
309
310 fn build_snapshot(self) -> Result<ToolPackToolSnapshot, AgentError> {
311 let effect_class = self
312 .effect_class
313 .ok_or_else(|| AgentError::missing_required_field("tool.effect_class"))?;
314 let risk_class = self
315 .risk_class
316 .ok_or_else(|| AgentError::missing_required_field("tool.risk_class"))?;
317 let capability_id = self
318 .capability_id
319 .unwrap_or_else(|| CapabilityId::new(format!("cap.tool.{}", self.tool_name)));
320 let namespace = self
321 .namespace
322 .unwrap_or_else(|| CapabilityNamespace::new(format!("tool.{}", self.tool_name)));
323 Ok(ToolPackToolSnapshot {
324 capability_id,
325 canonical_tool_name: CanonicalToolName::new(self.tool_name),
326 namespace,
327 description: self.description,
328 schema_ref: PackageSidecarRef::new(self.schema_id, "tool_schema", "v1"),
329 redacted_schema: self.redacted_schema,
330 executor_ref: ExecutorRef::new(self.executor_ref),
331 policy_refs: self.policy_refs,
332 requires_approval: self.requires_approval,
333 required_permissions: self.required_permissions,
334 effect_class,
335 risk_class,
336 redaction_policy_ref: self.redaction_policy_ref,
337 timeout_ms: self.timeout_ms,
338 cancellation: self.cancellation,
339 reconciliation: self.reconciliation,
340 privacy: self.privacy,
341 })
342 }
343}
344
345#[derive(Clone, Debug)]
346pub struct ToolPackBuilder {
350 pack_id: ToolPackId,
351 kind: ToolPackKind,
352 version: String,
353 source: SourceRef,
354 trust: TrustClass,
355 tools: Vec<ToolPackToolSnapshot>,
356}
357
358impl ToolPackBuilder {
359 pub fn new(
361 pack_id: ToolPackId,
362 kind: ToolPackKind,
363 version: impl Into<String>,
364 source: SourceRef,
365 ) -> Self {
366 Self {
367 pack_id,
368 kind,
369 version: version.into(),
370 source,
371 trust: TrustClass::SdkGenerated,
372 tools: Vec::new(),
373 }
374 }
375
376 pub fn trust(mut self, trust: TrustClass) -> Self {
378 self.trust = trust;
379 self
380 }
381
382 pub fn listen(mut self, tool: Tool) -> Self {
385 self.tools.push(tool.into_snapshot());
386 self
387 }
388
389 pub fn listen_async(mut self, tool: AsyncTool) -> Self {
392 self.tools.push(tool.into_snapshot());
393 self
394 }
395
396 pub fn build(self) -> Result<ToolkitPackBundle, AgentError> {
399 let mut snapshot =
400 ToolPackSnapshot::new(self.pack_id, self.kind, self.version, self.source)
401 .with_trust(self.trust);
402 for tool in self.tools {
403 snapshot = snapshot.with_tool(tool);
404 }
405 ToolkitPackBundle::from_snapshot(snapshot)
406 }
407}