Skip to main content

a3s_code_core/tools/
registry.rs

1//! Tool Registry
2//!
3//! Central registry for all tools (built-in and dynamic).
4//! Provides thread-safe registration, lookup, and execution.
5
6use super::artifacts::{ArtifactStore, ArtifactStoreLimits, ToolArtifact};
7use super::types::{Tool, ToolCapabilities, ToolContext, ToolOutput};
8use super::ToolResult;
9use super::{
10    merge_tool_output_artifact_metadata, tool_output_artifact, transform_tool_output_with_artifact,
11    ImmutableContentAdapterSession, ImmutableContentError, ImmutableContentKindV1,
12    ToolOutputArtifact, ToolResultTransformBindingV1, ToolResultTransformPolicyV1,
13    TOOL_RESULT_CONTENT_MEDIA_TYPE,
14};
15use crate::llm::ToolDefinition;
16use crate::trace::{InMemoryTraceSink, TraceEvent, TraceSink};
17use anyhow::Result;
18use std::collections::HashMap;
19use std::hash::{Hash, Hasher};
20use std::path::PathBuf;
21use std::sync::{Arc, RwLock};
22use thiserror::Error;
23
24const MAX_TOOL_SCHEMA_BYTES: usize = 256 * 1024;
25const MAX_ARGUMENT_VALIDATION_ERRORS: usize = 8;
26const MAX_ARGUMENT_VALIDATION_MESSAGE_BYTES: usize = 4 * 1024;
27const MAX_INLINE_CHANGE_BYTES: usize = 64 * 1024;
28const CHANGE_SIDE_PREVIEW_BYTES: usize = 8 * 1024;
29const CHANGE_DIFF_PREVIEW_BYTES: usize = 32 * 1024;
30const MAX_DIFF_COMPUTE_BYTES: usize = 1024 * 1024;
31
32#[derive(Debug, Error)]
33#[error("projected tool name '{name}' conflicts with the compatibility registry")]
34pub(crate) struct ToolRegistrySnapshotError {
35    name: String,
36}
37
38impl ToolRegistrySnapshotError {
39    pub(crate) fn name(&self) -> &str {
40        &self.name
41    }
42}
43
44#[derive(Clone)]
45enum CachedArgumentValidator {
46    Valid(Arc<jsonschema::Validator>),
47    Invalid(String),
48}
49
50#[derive(Clone)]
51struct ArgumentValidatorCacheEntry {
52    schema_fingerprint: u64,
53    validator: CachedArgumentValidator,
54}
55
56/// Tool registry for managing all available tools
57pub struct ToolRegistry {
58    tools: RwLock<HashMap<String, Arc<dyn Tool>>>,
59    /// Names of builtin tools that cannot be overridden
60    builtins: RwLock<std::collections::HashSet<String>>,
61    context: RwLock<ToolContext>,
62    artifact_store: ArtifactStore,
63    immutable_content_adapter: Option<ImmutableContentAdapterSession>,
64    trace_sink: RwLock<Arc<dyn TraceSink>>,
65    argument_validators: RwLock<HashMap<String, ArgumentValidatorCacheEntry>>,
66    transform_policy: RwLock<ToolResultTransformPolicyV1>,
67}
68
69/// Output prepared by the registry's single observation pipeline.
70///
71/// Keeping the transform binding alongside the output prevents the `ToolResult`
72/// and `ToolOutput` entry points from independently reimplementing artifact,
73/// compaction, and evidence handling. The public raw-output API is intentionally
74/// only a shape projection over this prepared result.
75struct PreparedToolOutput {
76    output: Option<ToolOutput>,
77    transform_binding: ToolResultTransformBindingV1,
78}
79
80impl ToolRegistry {
81    /// Create a new tool registry
82    pub fn new(workspace: PathBuf) -> Self {
83        Self::with_artifact_limits(workspace, ArtifactStoreLimits::default())
84    }
85
86    /// Create a new tool registry with custom artifact retention limits.
87    pub fn with_artifact_limits(workspace: PathBuf, artifact_limits: ArtifactStoreLimits) -> Self {
88        Self::with_artifact_limits_and_workspace_services(
89            workspace.clone(),
90            artifact_limits,
91            crate::workspace::WorkspaceServices::local(workspace),
92        )
93    }
94
95    /// Create a new tool registry with custom artifact limits and workspace backend.
96    pub fn with_artifact_limits_and_workspace_services(
97        workspace: PathBuf,
98        artifact_limits: ArtifactStoreLimits,
99        workspace_services: Arc<crate::workspace::WorkspaceServices>,
100    ) -> Self {
101        Self::with_workspace_services_artifact_limits_and_immutable_content_adapter(
102            workspace,
103            artifact_limits,
104            workspace_services,
105            None,
106        )
107    }
108
109    pub(crate) fn with_workspace_services_artifact_limits_and_immutable_content_adapter(
110        workspace: PathBuf,
111        artifact_limits: ArtifactStoreLimits,
112        workspace_services: Arc<crate::workspace::WorkspaceServices>,
113        immutable_content_adapter: Option<ImmutableContentAdapterSession>,
114    ) -> Self {
115        let context = ToolContext::new(workspace).with_workspace_services(workspace_services);
116        Self {
117            tools: RwLock::new(HashMap::new()),
118            builtins: RwLock::new(std::collections::HashSet::new()),
119            context: RwLock::new(context),
120            artifact_store: ArtifactStore::with_limits(artifact_limits),
121            immutable_content_adapter,
122            trace_sink: RwLock::new(Arc::new(InMemoryTraceSink::default())),
123            argument_validators: RwLock::new(HashMap::new()),
124            transform_policy: RwLock::new(ToolResultTransformPolicyV1::default()),
125        }
126    }
127
128    /// Freeze the current compatibility registry and add one conflict-free
129    /// external generation. The returned registry shares service handles and
130    /// exact Tool `Arc`s but no mutable name map with the source registry.
131    pub(crate) fn snapshot_with_external_tools(
132        &self,
133        external: impl IntoIterator<Item = Arc<dyn Tool>>,
134    ) -> Result<Self, ToolRegistrySnapshotError> {
135        // Preserve the registry's established lock order: tools before
136        // builtins. The remaining fields have no operation that takes either
137        // lock while holding their own lock.
138        let mut tools = self.tools.read().unwrap().clone();
139        let builtins = self.builtins.read().unwrap().clone();
140        for tool in external {
141            let name = tool.name().to_owned();
142            if tools.contains_key(&name) {
143                return Err(ToolRegistrySnapshotError { name });
144            }
145            tools.insert(name, tool);
146        }
147        Ok(Self {
148            tools: RwLock::new(tools),
149            builtins: RwLock::new(builtins),
150            context: RwLock::new(self.context.read().unwrap().clone()),
151            artifact_store: self.artifact_store.clone(),
152            immutable_content_adapter: self.immutable_content_adapter.clone(),
153            trace_sink: RwLock::new(Arc::clone(&self.trace_sink.read().unwrap())),
154            argument_validators: RwLock::new(self.argument_validators.read().unwrap().clone()),
155            transform_policy: RwLock::new(self.transform_policy.read().unwrap().clone()),
156        })
157    }
158
159    pub(crate) fn set_tool_result_transform_policy(
160        &self,
161        policy: ToolResultTransformPolicyV1,
162    ) -> Result<()> {
163        policy.validate()?;
164        *self.transform_policy.write().unwrap() = policy;
165        Ok(())
166    }
167
168    /// Register a builtin tool (cannot be overridden by dynamic tools)
169    pub fn register_builtin(&self, tool: Arc<dyn Tool>) {
170        let name = tool.name().to_string();
171        let mut tools = self.tools.write().unwrap();
172        let mut builtins = self.builtins.write().unwrap();
173        tracing::debug!("Registering builtin tool: {}", name);
174        tools.insert(name.clone(), tool);
175        builtins.insert(name);
176    }
177
178    /// Register a tool
179    ///
180    /// If a tool with the same name already exists as a builtin, the registration
181    /// is rejected to prevent shadowing of core tools.
182    pub fn register(&self, tool: Arc<dyn Tool>) {
183        let name = tool.name().to_string();
184        // All operations that need both registry locks take `tools` first.
185        // This keeps the builtin check and insertion atomic with
186        // `register_builtin` and avoids lock-order inversion.
187        let mut tools = self.tools.write().unwrap();
188        let builtins = self.builtins.read().unwrap();
189        if builtins.contains(&name) {
190            tracing::warn!(
191                "Rejected registration of tool '{}': cannot shadow builtin",
192                name
193            );
194            return;
195        }
196        tracing::debug!("Registering tool: {}", name);
197        tools.insert(name, tool);
198    }
199
200    /// Register a dynamic tool and return the tool it shadowed.
201    ///
202    /// The lookup and replacement happen under one write lock so lifecycle
203    /// owners can later restore the exact prior registration without racing a
204    /// concurrent dynamic registration. The boolean is `false` when a builtin
205    /// owns the name and the dynamic registration was rejected.
206    pub(crate) fn register_with_shadow(
207        &self,
208        tool: Arc<dyn Tool>,
209    ) -> (bool, Option<Arc<dyn Tool>>) {
210        let name = tool.name().to_string();
211        let mut tools = self.tools.write().unwrap();
212        let builtins = self.builtins.read().unwrap();
213        if builtins.contains(&name) {
214            tracing::warn!(
215                "Rejected registration of tool '{}': cannot shadow builtin",
216                name
217            );
218            return (false, None);
219        }
220        tracing::debug!("Registering owned dynamic tool: {}", name);
221        (true, tools.insert(name, tool))
222    }
223
224    /// Restore a shadowed registration only while `expected` still owns the
225    /// name.
226    ///
227    /// This compare-and-replace prevents one lifecycle owner from deleting or
228    /// overwriting a tool installed later by another dynamic source.
229    pub(crate) fn restore_if_same(
230        &self,
231        name: &str,
232        expected: &Arc<dyn Tool>,
233        replacement: Option<Arc<dyn Tool>>,
234    ) -> bool {
235        let mut tools = self.tools.write().unwrap();
236        let Some(current) = tools.get(name) else {
237            return false;
238        };
239        if !Arc::ptr_eq(current, expected) {
240            return false;
241        }
242
243        match replacement {
244            Some(tool) => {
245                tools.insert(name.to_string(), tool);
246            }
247            None => {
248                tools.remove(name);
249            }
250        }
251        true
252    }
253
254    /// Register a dynamic tool only when no source currently owns its name.
255    pub(crate) fn register_if_absent(&self, tool: Arc<dyn Tool>) -> bool {
256        let name = tool.name().to_string();
257        let mut tools = self.tools.write().unwrap();
258        if tools.contains_key(&name) {
259            return false;
260        }
261        tracing::debug!("Registering previously absent dynamic tool: {}", name);
262        tools.insert(name, tool);
263        true
264    }
265
266    /// Unregister a tool by name
267    ///
268    /// Returns true if the tool was found and removed.
269    pub fn unregister(&self, name: &str) -> bool {
270        let mut tools = self.tools.write().unwrap();
271        let builtins = self.builtins.read().unwrap();
272        if builtins.contains(name) {
273            tracing::warn!(
274                "Rejected unregister of tool '{}': builtin tools cannot be removed through dynamic unregister",
275                name
276            );
277            return false;
278        }
279        tracing::debug!("Unregistering tool: {}", name);
280        tools.remove(name).is_some()
281    }
282
283    /// Unregister all tools whose names start with the given prefix.
284    pub fn unregister_by_prefix(&self, prefix: &str) {
285        let mut tools = self.tools.write().unwrap();
286        let builtins = self.builtins.read().unwrap();
287        tools.retain(|name, _| builtins.contains(name) || !name.starts_with(prefix));
288        tracing::debug!("Unregistered tools with prefix: {}", prefix);
289    }
290
291    /// Get a tool by name
292    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
293        let tools = self.tools.read().unwrap();
294        tools.get(name).cloned()
295    }
296
297    pub(crate) fn capabilities(
298        &self,
299        name: &str,
300        args: &serde_json::Value,
301    ) -> Option<ToolCapabilities> {
302        self.get(name).map(|tool| tool.capabilities(args))
303    }
304
305    pub(crate) fn requires_confirmation(&self, name: &str, args: &serde_json::Value) -> bool {
306        self.get(name)
307            .is_some_and(|tool| tool.requires_confirmation(args))
308    }
309
310    /// Check if a tool exists
311    pub fn contains(&self, name: &str) -> bool {
312        let tools = self.tools.read().unwrap();
313        tools.contains_key(name)
314    }
315
316    /// Get all tool definitions for LLM
317    pub fn definitions(&self) -> Vec<ToolDefinition> {
318        let tools = self.tools.read().unwrap();
319        let mut definitions = tools
320            .values()
321            .filter(|tool| tool.is_model_visible())
322            .map(|tool| tool.definition())
323            .collect::<Vec<_>>();
324        definitions.sort_by(|a, b| a.name.cmp(&b.name));
325        definitions
326    }
327
328    /// List all registered tool names
329    pub fn list(&self) -> Vec<String> {
330        let tools = self.tools.read().unwrap();
331        let mut names = tools.keys().cloned().collect::<Vec<_>>();
332        names.sort();
333        names
334    }
335
336    /// Validate model- or orchestrator-supplied arguments against the tool's
337    /// declared JSON Schema before permissions or execution side effects.
338    ///
339    /// Low-level standalone registry calls remain compatibility-oriented and
340    /// do not invoke this automatically. The governed agent/session gateway is
341    /// the enforcement boundary.
342    pub(crate) fn validate_arguments(
343        &self,
344        name: &str,
345        args: &serde_json::Value,
346    ) -> std::result::Result<(), String> {
347        let Some(tool) = self.get(name) else {
348            return Ok(());
349        };
350        let schema = tool.parameters();
351        let schema_bytes = serde_json::to_vec(&schema)
352            .map_err(|error| format!("tool parameter schema is not serializable: {error}"))?;
353        if schema_bytes.len() > MAX_TOOL_SCHEMA_BYTES {
354            return Err(format!(
355                "tool parameter schema exceeds the {} byte safety limit",
356                MAX_TOOL_SCHEMA_BYTES
357            ));
358        }
359        let mut hasher = std::collections::hash_map::DefaultHasher::new();
360        schema_bytes.hash(&mut hasher);
361        let schema_fingerprint = hasher.finish();
362        let cached = self
363            .argument_validators
364            .read()
365            .unwrap()
366            .get(name)
367            .filter(|entry| entry.schema_fingerprint == schema_fingerprint)
368            .cloned();
369        let validator = match cached.map(|entry| entry.validator) {
370            Some(CachedArgumentValidator::Valid(validator)) => validator,
371            Some(CachedArgumentValidator::Invalid(error)) => return Err(error),
372            None => {
373                let compiled = match jsonschema::draft202012::options().build(&schema) {
374                    Ok(validator) => CachedArgumentValidator::Valid(Arc::new(validator)),
375                    Err(error) => CachedArgumentValidator::Invalid(format!(
376                        "tool has an invalid parameter schema: {error}"
377                    )),
378                };
379                self.argument_validators.write().unwrap().insert(
380                    name.to_string(),
381                    ArgumentValidatorCacheEntry {
382                        schema_fingerprint,
383                        validator: compiled.clone(),
384                    },
385                );
386                match compiled {
387                    CachedArgumentValidator::Valid(validator) => validator,
388                    CachedArgumentValidator::Invalid(error) => return Err(error),
389                }
390            }
391        };
392        let mut errors = validator
393            .iter_errors(args)
394            .take(MAX_ARGUMENT_VALIDATION_ERRORS + 1)
395            .map(|error| {
396                let path = error.instance_path().to_string();
397                if path.is_empty() {
398                    format!("$: {error}")
399                } else {
400                    format!("{path}: {error}")
401                }
402            })
403            .collect::<Vec<_>>();
404        if errors.is_empty() {
405            return Ok(());
406        }
407
408        let omitted = errors.len() > MAX_ARGUMENT_VALIDATION_ERRORS;
409        errors.truncate(MAX_ARGUMENT_VALIDATION_ERRORS);
410        let mut message = errors.join("; ");
411        if omitted {
412            message.push_str("; additional validation errors omitted");
413        }
414        Err(crate::text::truncate_utf8(&message, MAX_ARGUMENT_VALIDATION_MESSAGE_BYTES).to_string())
415    }
416
417    /// Get the number of registered tools
418    pub fn len(&self) -> usize {
419        let tools = self.tools.read().unwrap();
420        tools.len()
421    }
422
423    /// Check if registry is empty
424    pub fn is_empty(&self) -> bool {
425        self.len() == 0
426    }
427
428    /// Get the tool context
429    pub fn context(&self) -> ToolContext {
430        self.context.read().unwrap().clone()
431    }
432
433    /// Return a clone of the registry's artifact store handle.
434    pub fn artifact_store(&self) -> ArtifactStore {
435        self.artifact_store.clone()
436    }
437
438    /// Return the session-scoped immutable-content adapter, when configured.
439    pub fn immutable_content_adapter(&self) -> Option<&ImmutableContentAdapterSession> {
440        self.immutable_content_adapter.as_ref()
441    }
442
443    /// Get a stored tool artifact by URI.
444    pub fn get_artifact(&self, artifact_uri: &str) -> Option<ToolArtifact> {
445        self.artifact_store.get(artifact_uri)
446    }
447
448    /// Replace the trace sink used for compact tool/program execution events.
449    pub fn set_trace_sink(&self, sink: Arc<dyn TraceSink>) {
450        *self.trace_sink.write().unwrap() = sink;
451    }
452
453    /// Return the current trace sink.
454    pub fn trace_sink(&self) -> Arc<dyn TraceSink> {
455        Arc::clone(&self.trace_sink.read().unwrap())
456    }
457
458    /// Set the search configuration for the tool context
459    pub fn set_search_config(&self, config: crate::config::SearchConfig) {
460        let mut ctx = self.context.write().unwrap();
461        *ctx = ctx.clone().with_search_config(config);
462    }
463
464    /// Set a sandbox executor so that `bash` tool calls use the sandbox even
465    /// when executed without an explicit `ToolContext` (i.e., via `execute()`).
466    pub fn set_sandbox(&self, sandbox: std::sync::Arc<dyn crate::sandbox::BashSandbox>) {
467        let mut ctx = self.context.write().unwrap();
468        *ctx = ctx.clone().with_sandbox(sandbox);
469    }
470
471    /// Set environment overrides used by subprocess-backed tools when executed
472    /// without an explicit context.
473    pub fn set_command_env(&self, env: Arc<HashMap<String, String>>) {
474        let mut ctx = self.context.write().unwrap();
475        *ctx = ctx.clone().with_command_env(env);
476    }
477
478    /// Execute a tool by name using the registry's default context.
479    ///
480    /// This is the lowest-level standalone registry boundary. It does not run
481    /// agent/session permission, HITL, hook, budget, queue, timeout,
482    /// cancellation, or sanitization policy.
483    pub async fn execute(&self, name: &str, args: &serde_json::Value) -> Result<ToolResult> {
484        let ctx = self.context();
485        self.execute_with_context(name, args, &ctx).await
486    }
487
488    /// Execute a tool by name with an external caller-owned context.
489    ///
490    /// This remains a low-level ungoverned call; agent/session paths must use
491    /// their scoped invocation gateway instead.
492    pub async fn execute_with_context(
493        &self,
494        name: &str,
495        args: &serde_json::Value,
496        ctx: &ToolContext,
497    ) -> Result<ToolResult> {
498        let start = std::time::Instant::now();
499        let prepared = self.prepare_output_with_context(name, args, ctx).await?;
500        let transform_binding = prepared.transform_binding;
501        let mut result = match prepared.output {
502            Some(output) => Ok(ToolResult {
503                name: name.to_string(),
504                output: output.content,
505                exit_code: if output.success { 0 } else { 1 },
506                metadata: output.metadata,
507                images: output.images,
508                error_kind: output.error_kind,
509                trust: crate::tools::ToolResultTrustV1::WorkspaceData,
510            }),
511            None => Ok(ToolResult::error(name, format!("Unknown tool: {name}"))),
512        };
513
514        if let Ok(result) = &mut result {
515            result.metadata = Some(super::ensure_tool_result_evidence_with_transform_binding(
516                result.metadata.take(),
517                &result.output,
518                &transform_binding,
519            )?);
520        }
521
522        if let Ok(ref r) = result {
523            crate::telemetry::record_tool_result(r.exit_code, start.elapsed());
524            self.record_trace_event(name, r, start.elapsed());
525        }
526
527        result
528    }
529
530    /// Execute a tool and return the `ToolOutput` projection using the registry's
531    /// default context. The same observation pipeline as `execute` is applied.
532    pub async fn execute_raw(
533        &self,
534        name: &str,
535        args: &serde_json::Value,
536    ) -> Result<Option<ToolOutput>> {
537        let ctx = self.context();
538        self.execute_raw_with_context(name, args, &ctx).await
539    }
540
541    /// Execute a tool and return the `ToolOutput` projection with an external
542    /// context. The same observation pipeline as `execute_with_context` is applied.
543    pub async fn execute_raw_with_context(
544        &self,
545        name: &str,
546        args: &serde_json::Value,
547        ctx: &ToolContext,
548    ) -> Result<Option<ToolOutput>> {
549        Ok(self
550            .prepare_output_with_context(name, args, ctx)
551            .await?
552            .output)
553    }
554
555    async fn prepare_output_with_context(
556        &self,
557        name: &str,
558        args: &serde_json::Value,
559        ctx: &ToolContext,
560    ) -> Result<PreparedToolOutput> {
561        let policy = self.transform_policy.read().unwrap().clone();
562        let transform_binding = ToolResultTransformBindingV1::from_policy(&policy)?;
563        let Some(tool) = self.get(name) else {
564            return Ok(PreparedToolOutput {
565                output: None,
566                transform_binding,
567            });
568        };
569
570        let mut output = tool.execute(args, ctx).await?;
571        self.compact_change_metadata(name, &mut output.metadata, ctx)
572            .await?;
573        let original_content = output.content.clone();
574        let truncated = transform_tool_output_with_artifact(name, &output.content, &policy);
575        output.content = truncated.content;
576        let loss_mode = truncated.loss_mode;
577        let projected_artifact_reference = truncated.artifact.is_some();
578        let artifact = truncated.artifact.or_else(|| {
579            self.immutable_content_adapter
580                .as_ref()
581                .map(|_| super::tool_output_artifact(name, &original_content, output.content.len()))
582        });
583        if let Some(mut artifact) = artifact {
584            let compatibility_uri = artifact.artifact_uri.clone();
585            self.store_tool_artifact(
586                name,
587                &original_content,
588                &mut artifact,
589                ImmutableContentKindV1::ToolResultOriginal,
590                ctx,
591            )
592            .await?;
593            if projected_artifact_reference {
594                rewrite_projected_artifact_uri(
595                    &mut output.content,
596                    &compatibility_uri,
597                    &artifact.artifact_uri,
598                )?;
599            }
600            output.metadata = Some(merge_tool_output_artifact_metadata(
601                output.metadata,
602                &artifact,
603            ));
604        }
605        output.metadata = Some(super::attach_tool_result_evidence_with_transform_binding(
606            output.metadata,
607            &original_content,
608            &output.content,
609            loss_mode,
610            &transform_binding,
611        )?);
612        Ok(PreparedToolOutput {
613            output: Some(output),
614            transform_binding,
615        })
616    }
617
618    async fn store_tool_artifact(
619        &self,
620        tool_name: &str,
621        content: &str,
622        artifact: &mut ToolOutputArtifact,
623        kind: ImmutableContentKindV1,
624        ctx: &ToolContext,
625    ) -> Result<()> {
626        if let Some(adapter) = &self.immutable_content_adapter {
627            let cancellation = ctx.cancellation_token();
628            let retained = tokio::select! {
629                biased;
630                _ = cancellation.cancelled() => {
631                    return Err(ImmutableContentError::Cancelled.into());
632                }
633                result = adapter.put(kind, TOOL_RESULT_CONTENT_MEDIA_TYPE, content.as_bytes()) => {
634                    result.map_err(|error| anyhow::anyhow!(
635                        "immutable content adapter '{}' rejected Tool content: {}",
636                        adapter.adapter_name(),
637                        error.redacted_message(),
638                    ))?
639                }
640            };
641            artifact.artifact_uri.clone_from(&retained.uri);
642            artifact.content_reference = Some(retained);
643            return Ok(());
644        }
645
646        self.artifact_store
647            .put_content_addressed(ToolArtifact {
648                artifact_id: artifact.artifact_id.clone(),
649                artifact_uri: artifact.artifact_uri.clone(),
650                tool_name: tool_name.to_string(),
651                content: content.to_string(),
652                original_bytes: artifact.original_bytes,
653                shown_bytes: artifact.shown_bytes,
654            })
655            .map_err(|error| anyhow::anyhow!("immutable Tool artifact conflict: {error}"))?;
656        Ok(())
657    }
658
659    async fn compact_change_metadata(
660        &self,
661        tool_name: &str,
662        metadata: &mut Option<serde_json::Value>,
663        ctx: &ToolContext,
664    ) -> Result<()> {
665        let Some(serde_json::Value::Object(object)) = metadata.as_mut() else {
666            return Ok(());
667        };
668        let before = object
669            .get("before")
670            .and_then(serde_json::Value::as_str)
671            .map(ToString::to_string);
672        let after = object
673            .get("after")
674            .and_then(serde_json::Value::as_str)
675            .map(ToString::to_string);
676        if before.is_none() && after.is_none() {
677            return Ok(());
678        }
679
680        let before_bytes = before.as_ref().map_or(0, String::len);
681        let after_bytes = after.as_ref().map_or(0, String::len);
682        let total_bytes = before_bytes.saturating_add(after_bytes);
683        let compacted = total_bytes > MAX_INLINE_CHANGE_BYTES;
684        let before_artifact = match before.as_deref() {
685            Some(content) => {
686                self.store_change_artifact(
687                    tool_name,
688                    "before",
689                    content,
690                    compacted,
691                    ImmutableContentKindV1::ToolChangeBefore,
692                    ctx,
693                )
694                .await?
695            }
696            None => None,
697        };
698        let after_artifact = match after.as_deref() {
699            Some(content) => {
700                self.store_change_artifact(
701                    tool_name,
702                    "after",
703                    content,
704                    compacted,
705                    ImmutableContentKindV1::ToolChangeAfter,
706                    ctx,
707                )
708                .await?
709            }
710            None => None,
711        };
712
713        let unified_diff = if compacted && total_bytes <= MAX_DIFF_COMPUTE_BYTES {
714            let diff = similar::TextDiff::from_lines(
715                before.as_deref().unwrap_or_default(),
716                after.as_deref().unwrap_or_default(),
717            )
718            .unified_diff()
719            .context_radius(3)
720            .header("before", "after")
721            .to_string();
722            Some(bounded_head_tail(&diff, CHANGE_DIFF_PREVIEW_BYTES))
723        } else {
724            None
725        };
726
727        if compacted {
728            if let Some(content) = before.as_deref() {
729                object.insert(
730                    "before".to_string(),
731                    serde_json::Value::String(bounded_head_tail(
732                        content,
733                        CHANGE_SIDE_PREVIEW_BYTES,
734                    )),
735                );
736            }
737            if let Some(content) = after.as_deref() {
738                object.insert(
739                    "after".to_string(),
740                    serde_json::Value::String(bounded_head_tail(
741                        content,
742                        CHANGE_SIDE_PREVIEW_BYTES,
743                    )),
744                );
745            }
746        }
747
748        object.insert(
749            "change".to_string(),
750            serde_json::json!({
751                "compacted": compacted,
752                "before": before.as_deref().map(|content| serde_json::json!({
753                    "bytes": content.len(),
754                    "sha256": sha256::digest(content.as_bytes()),
755                    "artifact": before_artifact,
756                })),
757                "after": after.as_deref().map(|content| serde_json::json!({
758                    "bytes": content.len(),
759                    "sha256": sha256::digest(content.as_bytes()),
760                    "artifact": after_artifact,
761                })),
762                "unified_diff": unified_diff,
763                "diff_omitted": compacted && total_bytes > MAX_DIFF_COMPUTE_BYTES,
764            }),
765        );
766        Ok(())
767    }
768
769    async fn store_change_artifact(
770        &self,
771        tool_name: &str,
772        side: &str,
773        content: &str,
774        store: bool,
775        kind: ImmutableContentKindV1,
776        ctx: &ToolContext,
777    ) -> Result<Option<serde_json::Value>> {
778        if !store
779            || self.immutable_content_adapter.is_none()
780                && content.len() > self.artifact_store.limits().max_bytes
781        {
782            return Ok(None);
783        }
784        let mut artifact = tool_output_artifact(&format!("{tool_name}-{side}"), content, 0);
785        self.store_tool_artifact(tool_name, content, &mut artifact, kind, ctx)
786            .await?;
787        let mut metadata = serde_json::json!({
788            "artifact_id": artifact.artifact_id,
789            "artifact_uri": artifact.artifact_uri,
790        });
791        if let Some(reference) = artifact.content_reference {
792            metadata["content_reference"] = serde_json::json!(reference);
793        }
794        Ok(Some(metadata))
795    }
796
797    fn record_trace_event(&self, name: &str, result: &ToolResult, duration: std::time::Duration) {
798        let sink = self.trace_sink();
799        sink.record(TraceEvent::tool_execution(
800            name,
801            result.exit_code == 0,
802            result.exit_code,
803            duration,
804            result.output.len(),
805            result.metadata.as_ref(),
806        ));
807
808        if name == "program" {
809            sink.record(TraceEvent::program_execution(
810                name,
811                result.exit_code == 0,
812                result.exit_code,
813                duration,
814                result.output.len(),
815                result.metadata.as_ref(),
816            ));
817        }
818    }
819}
820
821fn rewrite_projected_artifact_uri(
822    projected: &mut String,
823    compatibility_uri: &str,
824    retained_uri: &str,
825) -> Result<()> {
826    if compatibility_uri == retained_uri {
827        return Ok(());
828    }
829    let start = projected.rfind(compatibility_uri).ok_or_else(|| {
830        anyhow::anyhow!(
831            "Tool result projection lost its compatibility artifact reference before retention"
832        )
833    })?;
834    let end = start + compatibility_uri.len();
835    projected.replace_range(start..end, retained_uri);
836    Ok(())
837}
838
839fn bounded_head_tail(content: &str, max_bytes: usize) -> String {
840    if content.len() <= max_bytes {
841        return content.to_string();
842    }
843    let head_limit = max_bytes / 2;
844    let tail_limit = max_bytes.saturating_sub(head_limit);
845    let head = crate::text::truncate_utf8(content, head_limit);
846    let mut tail_start = content.len().saturating_sub(tail_limit);
847    while tail_start < content.len() && !content.is_char_boundary(tail_start) {
848        tail_start += 1;
849    }
850    format!(
851        "{}\n\n... [{} bytes omitted from middle] ...\n\n{}",
852        head,
853        content
854            .len()
855            .saturating_sub(head.len())
856            .saturating_sub(content.len().saturating_sub(tail_start)),
857        &content[tail_start..]
858    )
859}
860
861#[cfg(test)]
862#[path = "registry/tests.rs"]
863mod tests;