Skip to main content

llm_tool_runtime/
starter_tools.rs

1use crate::{
2    Tool, ToolApprovalKind, ToolBackendKind, ToolCall, ToolCtx, ToolDescriptor, ToolError,
3    ToolExposureMode, ToolExposurePolicy, ToolIdempotencyClass, ToolOutputMode,
4    ToolReceiptPersistence, ToolResult, ToolSideEffectClass,
5};
6use async_trait::async_trait;
7use serde::{Deserialize, Serialize};
8use serde_json::json;
9use stack_ids::{ArtifactId, ScopeKey};
10use std::sync::Arc;
11
12#[async_trait]
13pub trait MemorySearchPort: Send + Sync {
14    async fn search(
15        &self,
16        query: &str,
17        scope: Option<&ScopeKey>,
18        limit: usize,
19    ) -> Result<serde_json::Value, ToolError>;
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct ArtifactContent {
24    pub artifact_id: ArtifactId,
25    pub content: String,
26}
27
28#[async_trait]
29pub trait ArtifactReader: Send + Sync {
30    async fn read(&self, artifact_id: &ArtifactId) -> Result<ArtifactContent, ToolError>;
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct VerificationJobRequest {
35    pub candidate_id: String,
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub task_id: Option<String>,
38}
39
40#[async_trait]
41pub trait RunVerificationPort: Send + Sync {
42    async fn enqueue_verification(
43        &self,
44        request: VerificationJobRequest,
45        dry_run: bool,
46        execution_permit: &crate::ToolExecutionPermit,
47    ) -> Result<crate::ToolJobHandle, ToolError>;
48}
49
50#[async_trait]
51pub trait ApplyPatchPreviewPort: Send + Sync {
52    async fn preview_patch(
53        &self,
54        target: &str,
55        patch: &str,
56        execution_permit: &crate::ToolExecutionPermit,
57    ) -> Result<serde_json::Value, ToolError>;
58}
59
60#[async_trait]
61pub trait SubmitPatchPort: Send + Sync {
62    async fn submit_patch(
63        &self,
64        target: &str,
65        patch: &str,
66        approval_grant: &crate::ApprovalGrant,
67        execution_permit: &crate::ToolExecutionPermit,
68    ) -> Result<serde_json::Value, ToolError>;
69}
70
71fn scoped_execution_permit<'a>(
72    ctx: &'a ToolCtx,
73    tool_name: &str,
74    target_key: &str,
75) -> Result<&'a crate::ToolExecutionPermit, ToolError> {
76    let permit = ctx.execution_permit.as_ref().ok_or_else(|| {
77        ToolError::new(
78            crate::ToolErrorClass::ApprovalRequired,
79            format!("{tool_name} requires an execution permit"),
80        )
81    })?;
82    let scope = ctx.scope.as_ref().ok_or_else(|| {
83        ToolError::new(
84            crate::ToolErrorClass::Denied,
85            format!("{tool_name} requires a scoped context"),
86        )
87    })?;
88    if permit.scope().namespace() != scope.namespace || permit.scope().target_key() != target_key {
89        return Err(ToolError::new(
90            crate::ToolErrorClass::Denied,
91            format!("{tool_name} execution permit does not cover target {target_key}"),
92        ));
93    }
94    Ok(permit)
95}
96
97pub struct SearchMemoryTool<P> {
98    port: Arc<P>,
99    descriptor: ToolDescriptor,
100}
101
102impl<P> SearchMemoryTool<P> {
103    /// Builds the default search-memory tool descriptor around the supplied port.
104    pub fn new(port: Arc<P>) -> Self {
105        Self {
106            port,
107            descriptor: ToolDescriptor {
108                name: "search_memory".into(),
109                version: "1.0.0".into(),
110                description: Some("Search projected memory within the active scope".into()),
111                backend_kind: ToolBackendKind::LocalFunction,
112                input_schema: json!({
113                    "type": "object",
114                    "required": ["query"],
115                    "properties": {
116                        "query": {"type": "string"},
117                        "limit": {"type": "integer"}
118                    },
119                    "additionalProperties": false
120                }),
121                output_mode: ToolOutputMode::StructuredJson,
122                read_only: true,
123                side_effect_class: ToolSideEffectClass::ReadOnly,
124                idempotency_class: ToolIdempotencyClass::Idempotent,
125                approval_kind: ToolApprovalKind::None,
126                timeout_ms: 10_000,
127                concurrency_key: Some("search_memory".into()),
128                cache_ttl_ms: Some(5_000),
129                exposure_mode: ToolExposureMode::Auto,
130                mcp_surface_kind: crate::McpSurfaceKind::Tool,
131                exposure_policy: ToolExposurePolicy::default(),
132                receipt_persistence: ToolReceiptPersistence::Ephemeral,
133                output_size_limit_bytes: Some(128 * 1024),
134                provider_payload: None,
135            },
136        }
137    }
138}
139
140#[async_trait]
141impl<P> Tool for SearchMemoryTool<P>
142where
143    P: MemorySearchPort + 'static,
144{
145    fn descriptor(&self) -> &ToolDescriptor {
146        &self.descriptor
147    }
148
149    async fn invoke(&self, ctx: &ToolCtx, call: &ToolCall) -> Result<ToolResult, ToolError> {
150        let query = call
151            .arguments
152            .get("query")
153            .and_then(|value| value.as_str())
154            .unwrap_or_default();
155        let limit = call
156            .arguments
157            .get("limit")
158            .and_then(|value| value.as_u64())
159            .unwrap_or(5) as usize;
160        let results = self.port.search(query, ctx.scope.as_ref(), limit).await?;
161        Ok(ToolResult::json(results))
162    }
163}
164
165pub struct ReadArtifactTool<P> {
166    port: Arc<P>,
167    descriptor: ToolDescriptor,
168}
169
170impl<P> ReadArtifactTool<P> {
171    /// Builds the default read-artifact tool descriptor around the supplied port.
172    pub fn new(port: Arc<P>) -> Self {
173        Self {
174            port,
175            descriptor: ToolDescriptor {
176                name: "read_artifact".into(),
177                version: "1.0.0".into(),
178                description: Some("Read the content of a stored artifact".into()),
179                backend_kind: ToolBackendKind::LocalFunction,
180                input_schema: json!({
181                    "type": "object",
182                    "required": ["artifact_id"],
183                    "properties": {
184                        "artifact_id": {"type": "string"}
185                    },
186                    "additionalProperties": false
187                }),
188                output_mode: ToolOutputMode::Text,
189                read_only: true,
190                side_effect_class: ToolSideEffectClass::ReadOnly,
191                idempotency_class: ToolIdempotencyClass::Idempotent,
192                approval_kind: ToolApprovalKind::None,
193                timeout_ms: 10_000,
194                concurrency_key: Some("read_artifact".into()),
195                cache_ttl_ms: Some(5_000),
196                exposure_mode: ToolExposureMode::OptIn,
197                mcp_surface_kind: crate::McpSurfaceKind::Resource,
198                exposure_policy: ToolExposurePolicy::default(),
199                receipt_persistence: ToolReceiptPersistence::Ephemeral,
200                output_size_limit_bytes: Some(256 * 1024),
201                provider_payload: None,
202            },
203        }
204    }
205}
206
207#[async_trait]
208impl<P> Tool for ReadArtifactTool<P>
209where
210    P: ArtifactReader + 'static,
211{
212    fn descriptor(&self) -> &ToolDescriptor {
213        &self.descriptor
214    }
215
216    async fn invoke(&self, _ctx: &ToolCtx, call: &ToolCall) -> Result<ToolResult, ToolError> {
217        let artifact_id = call
218            .arguments
219            .get("artifact_id")
220            .and_then(|value| value.as_str())
221            .unwrap_or_default();
222        let content = self.port.read(&ArtifactId::new(artifact_id)).await?;
223        Ok(ToolResult::text(content.content))
224    }
225}
226
227pub struct RunVerificationTool<P> {
228    port: Arc<P>,
229    descriptor: ToolDescriptor,
230}
231
232impl<P> RunVerificationTool<P> {
233    /// Builds the default run-verification tool descriptor around the supplied port.
234    pub fn new(port: Arc<P>) -> Self {
235        Self {
236            port,
237            descriptor: ToolDescriptor {
238                name: "run_verification".into(),
239                version: "1.0.0".into(),
240                description: Some(
241                    "Enqueue a verification job and return a durable job handle".into(),
242                ),
243                backend_kind: ToolBackendKind::LocalFunction,
244                input_schema: json!({
245                    "type": "object",
246                    "required": ["candidate_id"],
247                    "properties": {
248                        "candidate_id": {"type": "string"},
249                        "task_id": {"type": "string"}
250                    },
251                    "additionalProperties": false
252                }),
253                output_mode: ToolOutputMode::JobHandle,
254                read_only: false,
255                side_effect_class: ToolSideEffectClass::Analysis,
256                idempotency_class: ToolIdempotencyClass::BestEffort,
257                approval_kind: ToolApprovalKind::PolicyRequired,
258                timeout_ms: 30_000,
259                concurrency_key: Some("run_verification".into()),
260                cache_ttl_ms: None,
261                exposure_mode: ToolExposureMode::OptIn,
262                mcp_surface_kind: crate::McpSurfaceKind::Tool,
263                exposure_policy: ToolExposurePolicy::default(),
264                receipt_persistence: ToolReceiptPersistence::ForgeRaw,
265                output_size_limit_bytes: Some(8 * 1024),
266                provider_payload: None,
267            },
268        }
269    }
270}
271
272#[async_trait]
273impl<P> Tool for RunVerificationTool<P>
274where
275    P: RunVerificationPort + 'static,
276{
277    fn descriptor(&self) -> &ToolDescriptor {
278        &self.descriptor
279    }
280
281    async fn invoke(&self, ctx: &ToolCtx, call: &ToolCall) -> Result<ToolResult, ToolError> {
282        let request: VerificationJobRequest = serde_json::from_value(call.arguments.clone())
283            .map_err(|error| ToolError {
284                class: crate::ToolErrorClass::InvalidArguments,
285                message: error.to_string(),
286                retryable: false,
287                details: None,
288            })?;
289        let permit =
290            scoped_execution_permit(ctx, self.descriptor.name.as_str(), &request.candidate_id)?;
291        let handle = self
292            .port
293            .enqueue_verification(request, ctx.dry_run, permit)
294            .await?;
295        Ok(ToolResult::job_handle(handle))
296    }
297}
298
299pub struct ApplyPatchPreviewTool<P> {
300    port: Arc<P>,
301    descriptor: ToolDescriptor,
302}
303
304impl<P> ApplyPatchPreviewTool<P> {
305    /// Builds the default patch-preview tool descriptor around the supplied port.
306    pub fn new(port: Arc<P>) -> Self {
307        Self {
308            port,
309            descriptor: ToolDescriptor {
310                name: "apply_patch_preview".into(),
311                version: "1.0.0".into(),
312                description: Some("Preview a patch application without mutating the target".into()),
313                backend_kind: ToolBackendKind::LocalFunction,
314                input_schema: json!({
315                    "type": "object",
316                    "required": ["target", "patch"],
317                    "properties": {
318                        "target": {"type": "string"},
319                        "patch": {"type": "string"}
320                    },
321                    "additionalProperties": false
322                }),
323                output_mode: ToolOutputMode::StructuredJson,
324                read_only: false,
325                side_effect_class: ToolSideEffectClass::PreviewWrite,
326                idempotency_class: ToolIdempotencyClass::Idempotent,
327                approval_kind: ToolApprovalKind::PolicyRequired,
328                timeout_ms: 10_000,
329                concurrency_key: Some("apply_patch_preview".into()),
330                cache_ttl_ms: None,
331                exposure_mode: ToolExposureMode::OptIn,
332                mcp_surface_kind: crate::McpSurfaceKind::Tool,
333                exposure_policy: ToolExposurePolicy::default(),
334                receipt_persistence: ToolReceiptPersistence::ForgeRaw,
335                output_size_limit_bytes: Some(64 * 1024),
336                provider_payload: None,
337            },
338        }
339    }
340}
341
342#[async_trait]
343impl<P> Tool for ApplyPatchPreviewTool<P>
344where
345    P: ApplyPatchPreviewPort + 'static,
346{
347    fn descriptor(&self) -> &ToolDescriptor {
348        &self.descriptor
349    }
350
351    async fn invoke(&self, _ctx: &ToolCtx, call: &ToolCall) -> Result<ToolResult, ToolError> {
352        let target = call
353            .arguments
354            .get("target")
355            .and_then(|value| value.as_str())
356            .unwrap_or_default();
357        let patch = call
358            .arguments
359            .get("patch")
360            .and_then(|value| value.as_str())
361            .unwrap_or_default();
362        let permit = scoped_execution_permit(_ctx, self.descriptor.name.as_str(), target)?;
363        let preview = self.port.preview_patch(target, patch, permit).await?;
364        Ok(ToolResult::json(preview))
365    }
366}
367
368pub struct SubmitPatchTool<P> {
369    port: Arc<P>,
370    descriptor: ToolDescriptor,
371}
372
373impl<P> SubmitPatchTool<P> {
374    /// Builds the default patch-submission tool descriptor around the supplied port.
375    pub fn new(port: Arc<P>) -> Self {
376        Self {
377            port,
378            descriptor: ToolDescriptor {
379                name: "submit_patch".into(),
380                version: "1.0.0".into(),
381                description: Some("Submit a patch to the target with explicit approval".into()),
382                backend_kind: ToolBackendKind::LocalFunction,
383                input_schema: json!({
384                    "type": "object",
385                    "required": ["target", "patch"],
386                    "properties": {
387                        "target": {"type": "string"},
388                        "patch": {"type": "string"}
389                    },
390                    "additionalProperties": false
391                }),
392                output_mode: ToolOutputMode::StructuredJson,
393                read_only: false,
394                side_effect_class: ToolSideEffectClass::Write,
395                idempotency_class: ToolIdempotencyClass::NonIdempotent,
396                approval_kind: ToolApprovalKind::UserRequired,
397                timeout_ms: 10_000,
398                concurrency_key: Some("submit_patch".into()),
399                cache_ttl_ms: None,
400                exposure_mode: ToolExposureMode::Hidden,
401                mcp_surface_kind: crate::McpSurfaceKind::Tool,
402                exposure_policy: ToolExposurePolicy {
403                    exposure_mode: ToolExposureMode::Hidden,
404                    ..Default::default()
405                },
406                receipt_persistence: ToolReceiptPersistence::ForgeRaw,
407                output_size_limit_bytes: Some(64 * 1024),
408                provider_payload: None,
409            },
410        }
411    }
412}
413
414#[async_trait]
415impl<P> Tool for SubmitPatchTool<P>
416where
417    P: SubmitPatchPort + 'static,
418{
419    fn descriptor(&self) -> &ToolDescriptor {
420        &self.descriptor
421    }
422
423    async fn invoke(&self, ctx: &ToolCtx, call: &ToolCall) -> Result<ToolResult, ToolError> {
424        let target = call
425            .arguments
426            .get("target")
427            .and_then(|value| value.as_str())
428            .unwrap_or_default();
429        let patch = call
430            .arguments
431            .get("patch")
432            .and_then(|value| value.as_str())
433            .unwrap_or_default();
434        let permit = scoped_execution_permit(ctx, self.descriptor.name.as_str(), target)?;
435        let result = self
436            .port
437            .submit_patch(
438                target,
439                patch,
440                ctx.approval_grant.as_ref().ok_or_else(|| {
441                    ToolError::new(
442                        crate::ToolErrorClass::ApprovalRequired,
443                        "submit_patch requires an approval grant",
444                    )
445                })?,
446                permit,
447            )
448            .await?;
449        Ok(ToolResult::json(result))
450    }
451}
452
453pub struct StarterToolPorts<MS, AR, RV, PP, SP> {
454    pub memory_search: Arc<MS>,
455    pub artifact_reader: Arc<AR>,
456    pub verification: Arc<RV>,
457    pub patch_preview: Arc<PP>,
458    pub patch_submit: Arc<SP>,
459}
460
461#[cfg(test)]
462mod tests {
463    use super::*;
464    use crate::{ToolCall, ToolCtx, ToolOriginKind, ToolRetryOwner};
465    use serde_json::json;
466    use stack_ids::{ArtifactId, AttemptId, TraceCtx, TrialId};
467
468    struct TestMemorySearch;
469
470    #[async_trait]
471    impl MemorySearchPort for TestMemorySearch {
472        async fn search(
473            &self,
474            query: &str,
475            scope: Option<&ScopeKey>,
476            limit: usize,
477        ) -> Result<serde_json::Value, ToolError> {
478            Ok(json!({
479                "query": query,
480                "scope": scope.map(|scope| scope.namespace.clone()),
481                "limit": limit,
482            }))
483        }
484    }
485
486    struct TestArtifactReader;
487
488    #[async_trait]
489    impl ArtifactReader for TestArtifactReader {
490        async fn read(&self, artifact_id: &ArtifactId) -> Result<ArtifactContent, ToolError> {
491            Ok(ArtifactContent {
492                artifact_id: artifact_id.clone(),
493                content: format!("artifact:{}", artifact_id.as_str()),
494            })
495        }
496    }
497
498    struct TestVerificationPort;
499
500    #[async_trait]
501    impl RunVerificationPort for TestVerificationPort {
502        async fn enqueue_verification(
503            &self,
504            request: VerificationJobRequest,
505            dry_run: bool,
506            execution_permit: &crate::ToolExecutionPermit,
507        ) -> Result<crate::ToolJobHandle, ToolError> {
508            Ok(crate::ToolJobHandle {
509                job_id: format!(
510                    "job:{}:{}",
511                    request.candidate_id,
512                    execution_permit.execution_permit_id()
513                ),
514                status: Some(if dry_run { "dry_run" } else { "queued" }.into()),
515            })
516        }
517    }
518
519    struct TestPatchPreviewPort;
520
521    #[async_trait]
522    impl ApplyPatchPreviewPort for TestPatchPreviewPort {
523        async fn preview_patch(
524            &self,
525            target: &str,
526            patch: &str,
527            execution_permit: &crate::ToolExecutionPermit,
528        ) -> Result<serde_json::Value, ToolError> {
529            Ok(json!({
530                "target": target,
531                "patch_len": patch.len(),
532                "execution_permit_id": execution_permit.execution_permit_id(),
533                "preview": true,
534            }))
535        }
536    }
537
538    struct TestSubmitPatchPort;
539
540    #[async_trait]
541    impl SubmitPatchPort for TestSubmitPatchPort {
542        async fn submit_patch(
543            &self,
544            target: &str,
545            patch: &str,
546            approval_grant: &crate::ApprovalGrant,
547            execution_permit: &crate::ToolExecutionPermit,
548        ) -> Result<serde_json::Value, ToolError> {
549            Ok(json!({
550                "target": target,
551                "patch_len": patch.len(),
552                "approval_grant_id": approval_grant.grant_id,
553                "execution_permit_id": execution_permit.execution_permit_id(),
554                "submitted": true,
555            }))
556        }
557    }
558
559    fn execution_permit_for(target_key: &str) -> crate::ToolExecutionPermit {
560        crate::ToolExecutionPermit::new(
561            stack_ids::ExecutionPermitId::generate(),
562            stack_ids::PolicyDecisionId::generate(),
563            None,
564            "starter-tools-test",
565            target_key,
566        )
567    }
568
569    fn tool_ctx() -> ToolCtx {
570        ToolCtx {
571            trace_ctx: TraceCtx::generate(),
572            attempt_id: AttemptId::generate(),
573            trial_id: TrialId::generate(),
574            deadline: None,
575            workload_class: None,
576            budget_context: None,
577            scope: Some(ScopeKey::namespace_only("starter-tools-test")),
578            dry_run: true,
579            approval_grant: Some(crate::ApprovalGrant::new(
580                "policy-1",
581                crate::ApprovalGrantScope {
582                    namespace: "starter-tools-test".into(),
583                    target_key: Some("src/lib.rs".into()),
584                    tool_name: "submit_patch".into(),
585                    effect_class: crate::ApprovalGrantEffectClass::Write,
586                    planner_stage: Some(crate::ToolPlannerStage::Execution),
587                },
588                vec!["operator".into()],
589                "2026-03-12T00:00:00Z",
590                None,
591                crate::ApprovalGrantCitation {
592                    applicability_context_id: stack_ids::ApplicabilityContextId::new(
593                        "applicability-context-1",
594                    ),
595                    profile_set_id: stack_ids::ProfileSetId::new("profile-set-1"),
596                    composition_receipt_id: stack_ids::CompositionReceiptId::new(
597                        "composition-receipt-1",
598                    ),
599                    effective_constitution_id: stack_ids::EffectiveConstitutionId::new(
600                        "effective-constitution-1",
601                    ),
602                    compiled_obligation_set_id: stack_ids::CompiledObligationSetId::new(
603                        "compiled-obligation-set-1",
604                    ),
605                },
606            )),
607            execution_permit: Some(execution_permit_for("src/lib.rs")),
608            idempotency_key: None,
609            caller: "starter-tools-test".into(),
610            planner_stage: crate::ToolPlannerStage::Execution,
611            parent_receipt_id: None,
612            family_receipt_id: None,
613            replay_parent_receipt_id: None,
614            remote_oracle_lease_id: None,
615            remote_slice_result_id: None,
616            attestation_envelope_id: None,
617            cross_runtime_replay_ticket_id: None,
618            retry_owner: Some(ToolRetryOwner::LlmPipeline),
619        }
620    }
621
622    #[tokio::test]
623    async fn starter_tools_return_expected_results() {
624        let ctx = tool_ctx();
625
626        let search_tool = SearchMemoryTool::new(Arc::new(TestMemorySearch));
627        let search_result = search_tool
628            .invoke(
629                &ctx,
630                &ToolCall::new(
631                    "search_memory",
632                    "1.0.0",
633                    json!({"query": "tool runtime", "limit": 3}),
634                    ToolOriginKind::Test,
635                ),
636            )
637            .await
638            .unwrap();
639        assert_eq!(search_result.payload["query"], "tool runtime");
640        assert_eq!(search_result.payload["limit"], 3);
641
642        let read_tool = ReadArtifactTool::new(Arc::new(TestArtifactReader));
643        let read_result = read_tool
644            .invoke(
645                &ctx,
646                &ToolCall::new(
647                    "read_artifact",
648                    "1.0.0",
649                    json!({"artifact_id": "artifact-1"}),
650                    ToolOriginKind::Test,
651                ),
652            )
653            .await
654            .unwrap();
655        assert_eq!(
656            read_result.display_text.as_deref(),
657            Some("artifact:artifact-1")
658        );
659
660        let verification_tool = RunVerificationTool::new(Arc::new(TestVerificationPort));
661        let mut verification_ctx = tool_ctx();
662        verification_ctx.execution_permit = Some(execution_permit_for("cand-1"));
663        let verification_result = verification_tool
664            .invoke(
665                &verification_ctx,
666                &ToolCall::new(
667                    "run_verification",
668                    "1.0.0",
669                    json!({"candidate_id": "cand-1"}),
670                    ToolOriginKind::Test,
671                ),
672            )
673            .await
674            .unwrap();
675        assert!(verification_result.payload["job_id"]
676            .as_str()
677            .is_some_and(|job_id| job_id.starts_with("job:cand-1:")));
678
679        let preview_tool = ApplyPatchPreviewTool::new(Arc::new(TestPatchPreviewPort));
680        let preview_result = preview_tool
681            .invoke(
682                &ctx,
683                &ToolCall::new(
684                    "apply_patch_preview",
685                    "1.0.0",
686                    json!({"target": "src/lib.rs", "patch": "*** Begin Patch"}),
687                    ToolOriginKind::Test,
688                ),
689            )
690            .await
691            .unwrap();
692        assert_eq!(preview_result.payload["preview"], true);
693        assert!(preview_result.payload["execution_permit_id"].is_string());
694
695        let submit_tool = SubmitPatchTool::new(Arc::new(TestSubmitPatchPort));
696        let submit_result = submit_tool
697            .invoke(
698                &ctx,
699                &ToolCall::new(
700                    "submit_patch",
701                    "1.0.0",
702                    json!({"target": "src/lib.rs", "patch": "*** Begin Patch"}),
703                    ToolOriginKind::Test,
704                ),
705            )
706            .await
707            .unwrap();
708        assert_eq!(submit_result.payload["submitted"], true);
709        assert!(submit_result.payload["approval_grant_id"].is_string());
710        assert!(submit_result.payload["execution_permit_id"].is_string());
711    }
712}