a3s_code_core/tools/

registry.rs

//! Tool Registry
//!
//! Central registry for all tools (built-in and dynamic).
//! Provides thread-safe registration, lookup, and execution.

use super::artifacts::{ArtifactStore, ArtifactStoreLimits, ToolArtifact};
use super::types::{Tool, ToolCapabilities, ToolContext, ToolOutput};
use super::ToolResult;
use super::{
    merge_tool_output_artifact_metadata, tool_output_artifact, transform_tool_output_with_artifact,
    ImmutableContentAdapterSession, ImmutableContentError, ImmutableContentKindV1,
    ToolOutputArtifact, ToolResultTransformBindingV1, ToolResultTransformPolicyV1,
    TOOL_RESULT_CONTENT_MEDIA_TYPE,
};
use crate::llm::ToolDefinition;
use crate::trace::{InMemoryTraceSink, TraceEvent, TraceSink};
use anyhow::Result;
use std::collections::HashMap;
use std::hash::{Hash, Hasher};
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use thiserror::Error;

const MAX_TOOL_SCHEMA_BYTES: usize = 256 * 1024;
const MAX_ARGUMENT_VALIDATION_ERRORS: usize = 8;
const MAX_ARGUMENT_VALIDATION_MESSAGE_BYTES: usize = 4 * 1024;
const MAX_INLINE_CHANGE_BYTES: usize = 64 * 1024;
const CHANGE_SIDE_PREVIEW_BYTES: usize = 8 * 1024;
const CHANGE_DIFF_PREVIEW_BYTES: usize = 32 * 1024;
const MAX_DIFF_COMPUTE_BYTES: usize = 1024 * 1024;

#[derive(Debug, Error)]
#[error("projected tool name '{name}' conflicts with the compatibility registry")]
pub(crate) struct ToolRegistrySnapshotError {
    name: String,
}

impl ToolRegistrySnapshotError {
    pub(crate) fn name(&self) -> &str {
        &self.name
    }
}

#[derive(Clone)]
enum CachedArgumentValidator {
    Valid(Arc<jsonschema::Validator>),
    Invalid(String),
}

#[derive(Clone)]
struct ArgumentValidatorCacheEntry {
    schema_fingerprint: u64,
    validator: CachedArgumentValidator,
}

/// Tool registry for managing all available tools
pub struct ToolRegistry {
    tools: RwLock<HashMap<String, Arc<dyn Tool>>>,
    /// Names of builtin tools that cannot be overridden
    builtins: RwLock<std::collections::HashSet<String>>,
    context: RwLock<ToolContext>,
    artifact_store: ArtifactStore,
    immutable_content_adapter: Option<ImmutableContentAdapterSession>,
    trace_sink: RwLock<Arc<dyn TraceSink>>,
    argument_validators: RwLock<HashMap<String, ArgumentValidatorCacheEntry>>,
    transform_policy: RwLock<ToolResultTransformPolicyV1>,
}

/// Output prepared by the registry's single observation pipeline.
///
/// Keeping the transform binding alongside the output prevents the `ToolResult`
/// and `ToolOutput` entry points from independently reimplementing artifact,
/// compaction, and evidence handling. The public raw-output API is intentionally
/// only a shape projection over this prepared result.
struct PreparedToolOutput {
    output: Option<ToolOutput>,
    transform_binding: ToolResultTransformBindingV1,
}

impl ToolRegistry {
    /// Create a new tool registry
    pub fn new(workspace: PathBuf) -> Self {
        Self::with_artifact_limits(workspace, ArtifactStoreLimits::default())
    }

    /// Create a new tool registry with custom artifact retention limits.
    pub fn with_artifact_limits(workspace: PathBuf, artifact_limits: ArtifactStoreLimits) -> Self {
        Self::with_artifact_limits_and_workspace_services(
            workspace.clone(),
            artifact_limits,
            crate::workspace::WorkspaceServices::local(workspace),
        )
    }

    /// Create a new tool registry with custom artifact limits and workspace backend.
    pub fn with_artifact_limits_and_workspace_services(
        workspace: PathBuf,
        artifact_limits: ArtifactStoreLimits,
        workspace_services: Arc<crate::workspace::WorkspaceServices>,
    ) -> Self {
        Self::with_workspace_services_artifact_limits_and_immutable_content_adapter(
            workspace,
            artifact_limits,
            workspace_services,
            None,
        )
    }

    pub(crate) fn with_workspace_services_artifact_limits_and_immutable_content_adapter(
        workspace: PathBuf,
        artifact_limits: ArtifactStoreLimits,
        workspace_services: Arc<crate::workspace::WorkspaceServices>,
        immutable_content_adapter: Option<ImmutableContentAdapterSession>,
    ) -> Self {
        let context = ToolContext::new(workspace).with_workspace_services(workspace_services);
        Self {
            tools: RwLock::new(HashMap::new()),
            builtins: RwLock::new(std::collections::HashSet::new()),
            context: RwLock::new(context),
            artifact_store: ArtifactStore::with_limits(artifact_limits),
            immutable_content_adapter,
            trace_sink: RwLock::new(Arc::new(InMemoryTraceSink::default())),
            argument_validators: RwLock::new(HashMap::new()),
            transform_policy: RwLock::new(ToolResultTransformPolicyV1::default()),
        }
    }

    /// Freeze the current compatibility registry and add one conflict-free
    /// external generation. The returned registry shares service handles and
    /// exact Tool `Arc`s but no mutable name map with the source registry.
    pub(crate) fn snapshot_with_external_tools(
        &self,
        external: impl IntoIterator<Item = Arc<dyn Tool>>,
    ) -> Result<Self, ToolRegistrySnapshotError> {
        // Preserve the registry's established lock order: tools before
        // builtins. The remaining fields have no operation that takes either
        // lock while holding their own lock.
        let mut tools = self.tools.read().unwrap().clone();
        let builtins = self.builtins.read().unwrap().clone();
        for tool in external {
            let name = tool.name().to_owned();
            if tools.contains_key(&name) {
                return Err(ToolRegistrySnapshotError { name });
            }
            tools.insert(name, tool);
        }
        Ok(Self {
            tools: RwLock::new(tools),
            builtins: RwLock::new(builtins),
            context: RwLock::new(self.context.read().unwrap().clone()),
            artifact_store: self.artifact_store.clone(),
            immutable_content_adapter: self.immutable_content_adapter.clone(),
            trace_sink: RwLock::new(Arc::clone(&self.trace_sink.read().unwrap())),
            argument_validators: RwLock::new(self.argument_validators.read().unwrap().clone()),
            transform_policy: RwLock::new(self.transform_policy.read().unwrap().clone()),
        })
    }

    pub(crate) fn set_tool_result_transform_policy(
        &self,
        policy: ToolResultTransformPolicyV1,
    ) -> Result<()> {
        policy.validate()?;
        *self.transform_policy.write().unwrap() = policy;
        Ok(())
    }

    /// Register a builtin tool (cannot be overridden by dynamic tools)
    pub fn register_builtin(&self, tool: Arc<dyn Tool>) {
        let name = tool.name().to_string();
        let mut tools = self.tools.write().unwrap();
        let mut builtins = self.builtins.write().unwrap();
        tracing::debug!("Registering builtin tool: {}", name);
        tools.insert(name.clone(), tool);
        builtins.insert(name);
    }

    /// Register a tool
    ///
    /// If a tool with the same name already exists as a builtin, the registration
    /// is rejected to prevent shadowing of core tools.
    pub fn register(&self, tool: Arc<dyn Tool>) {
        let name = tool.name().to_string();
        // All operations that need both registry locks take `tools` first.
        // This keeps the builtin check and insertion atomic with
        // `register_builtin` and avoids lock-order inversion.
        let mut tools = self.tools.write().unwrap();
        let builtins = self.builtins.read().unwrap();
        if builtins.contains(&name) {
            tracing::warn!(
                "Rejected registration of tool '{}': cannot shadow builtin",
                name
            );
            return;
        }
        tracing::debug!("Registering tool: {}", name);
        tools.insert(name, tool);
    }

    /// Register a dynamic tool and return the tool it shadowed.
    ///
    /// The lookup and replacement happen under one write lock so lifecycle
    /// owners can later restore the exact prior registration without racing a
    /// concurrent dynamic registration. The boolean is `false` when a builtin
    /// owns the name and the dynamic registration was rejected.
    pub(crate) fn register_with_shadow(
        &self,
        tool: Arc<dyn Tool>,
    ) -> (bool, Option<Arc<dyn Tool>>) {
        let name = tool.name().to_string();
        let mut tools = self.tools.write().unwrap();
        let builtins = self.builtins.read().unwrap();
        if builtins.contains(&name) {
            tracing::warn!(
                "Rejected registration of tool '{}': cannot shadow builtin",
                name
            );
            return (false, None);
        }
        tracing::debug!("Registering owned dynamic tool: {}", name);
        (true, tools.insert(name, tool))
    }

    /// Restore a shadowed registration only while `expected` still owns the
    /// name.
    ///
    /// This compare-and-replace prevents one lifecycle owner from deleting or
    /// overwriting a tool installed later by another dynamic source.
    pub(crate) fn restore_if_same(
        &self,
        name: &str,
        expected: &Arc<dyn Tool>,
        replacement: Option<Arc<dyn Tool>>,
    ) -> bool {
        let mut tools = self.tools.write().unwrap();
        let Some(current) = tools.get(name) else {
            return false;
        };
        if !Arc::ptr_eq(current, expected) {
            return false;
        }

        match replacement {
            Some(tool) => {
                tools.insert(name.to_string(), tool);
            }
            None => {
                tools.remove(name);
            }
        }
        true
    }

    /// Register a dynamic tool only when no source currently owns its name.
    pub(crate) fn register_if_absent(&self, tool: Arc<dyn Tool>) -> bool {
        let name = tool.name().to_string();
        let mut tools = self.tools.write().unwrap();
        if tools.contains_key(&name) {
            return false;
        }
        tracing::debug!("Registering previously absent dynamic tool: {}", name);
        tools.insert(name, tool);
        true
    }

    /// Unregister a tool by name
    ///
    /// Returns true if the tool was found and removed.
    pub fn unregister(&self, name: &str) -> bool {
        let mut tools = self.tools.write().unwrap();
        let builtins = self.builtins.read().unwrap();
        if builtins.contains(name) {
            tracing::warn!(
                "Rejected unregister of tool '{}': builtin tools cannot be removed through dynamic unregister",
                name
            );
            return false;
        }
        tracing::debug!("Unregistering tool: {}", name);
        tools.remove(name).is_some()
    }

    /// Unregister all tools whose names start with the given prefix.
    pub fn unregister_by_prefix(&self, prefix: &str) {
        let mut tools = self.tools.write().unwrap();
        let builtins = self.builtins.read().unwrap();
        tools.retain(|name, _| builtins.contains(name) || !name.starts_with(prefix));
        tracing::debug!("Unregistered tools with prefix: {}", prefix);
    }

    /// Get a tool by name
    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
        let tools = self.tools.read().unwrap();
        let canonical = crate::tool_name::canonical_tool_name(name);
        if let Some(tool) = tools.get(canonical) {
            return Some(tool.clone());
        }
        let lower = canonical.to_ascii_lowercase();
        if lower != canonical {
            if let Some(tool) = tools.get(&lower) {
                return Some(tool.clone());
            }
        }
        tools
            .iter()
            .find(|(registered, _)| registered.eq_ignore_ascii_case(canonical))
            .map(|(_, tool)| tool.clone())
    }

    pub(crate) fn capabilities(
        &self,
        name: &str,
        args: &serde_json::Value,
    ) -> Option<ToolCapabilities> {
        self.get(name).map(|tool| tool.capabilities(args))
    }

    pub(crate) fn requires_confirmation(&self, name: &str, args: &serde_json::Value) -> bool {
        self.get(name)
            .is_some_and(|tool| tool.requires_confirmation(args))
    }

    /// Check if a tool exists
    pub fn contains(&self, name: &str) -> bool {
        self.get(name).is_some()
    }

    /// Get all tool definitions for LLM
    pub fn definitions(&self) -> Vec<ToolDefinition> {
        let tools = self.tools.read().unwrap();
        let mut definitions = tools
            .values()
            .filter(|tool| tool.is_model_visible())
            .map(|tool| tool.definition())
            .collect::<Vec<_>>();
        definitions.sort_by(|a, b| a.name.cmp(&b.name));
        definitions
    }

    /// List all registered tool names
    pub fn list(&self) -> Vec<String> {
        let tools = self.tools.read().unwrap();
        let mut names = tools.keys().cloned().collect::<Vec<_>>();
        names.sort();
        names
    }

    /// Validate model- or orchestrator-supplied arguments against the tool's
    /// declared JSON Schema before permissions or execution side effects.
    ///
    /// Low-level standalone registry calls remain compatibility-oriented and
    /// do not invoke this automatically. The governed agent/session gateway is
    /// the enforcement boundary.
    pub(crate) fn validate_arguments(
        &self,
        name: &str,
        args: &serde_json::Value,
    ) -> std::result::Result<(), String> {
        let Some(tool) = self.get(name) else {
            return Ok(());
        };
        let schema = tool.parameters();
        let schema_bytes = serde_json::to_vec(&schema)
            .map_err(|error| format!("tool parameter schema is not serializable: {error}"))?;
        if schema_bytes.len() > MAX_TOOL_SCHEMA_BYTES {
            return Err(format!(
                "tool parameter schema exceeds the {} byte safety limit",
                MAX_TOOL_SCHEMA_BYTES
            ));
        }
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        schema_bytes.hash(&mut hasher);
        let schema_fingerprint = hasher.finish();
        let cached = self
            .argument_validators
            .read()
            .unwrap()
            .get(name)
            .filter(|entry| entry.schema_fingerprint == schema_fingerprint)
            .cloned();
        let validator = match cached.map(|entry| entry.validator) {
            Some(CachedArgumentValidator::Valid(validator)) => validator,
            Some(CachedArgumentValidator::Invalid(error)) => return Err(error),
            None => {
                let compiled = match jsonschema::draft202012::options().build(&schema) {
                    Ok(validator) => CachedArgumentValidator::Valid(Arc::new(validator)),
                    Err(error) => CachedArgumentValidator::Invalid(format!(
                        "tool has an invalid parameter schema: {error}"
                    )),
                };
                self.argument_validators.write().unwrap().insert(
                    name.to_string(),
                    ArgumentValidatorCacheEntry {
                        schema_fingerprint,
                        validator: compiled.clone(),
                    },
                );
                match compiled {
                    CachedArgumentValidator::Valid(validator) => validator,
                    CachedArgumentValidator::Invalid(error) => return Err(error),
                }
            }
        };
        let mut errors = validator
            .iter_errors(args)
            .take(MAX_ARGUMENT_VALIDATION_ERRORS + 1)
            .map(|error| {
                let path = error.instance_path().to_string();
                if path.is_empty() {
                    format!("$: {error}")
                } else {
                    format!("{path}: {error}")
                }
            })
            .collect::<Vec<_>>();
        if errors.is_empty() {
            return Ok(());
        }

        let omitted = errors.len() > MAX_ARGUMENT_VALIDATION_ERRORS;
        errors.truncate(MAX_ARGUMENT_VALIDATION_ERRORS);
        let mut message = errors.join("; ");
        if omitted {
            message.push_str("; additional validation errors omitted");
        }
        Err(crate::text::truncate_utf8(&message, MAX_ARGUMENT_VALIDATION_MESSAGE_BYTES).to_string())
    }

    /// Get the number of registered tools
    pub fn len(&self) -> usize {
        let tools = self.tools.read().unwrap();
        tools.len()
    }

    /// Check if registry is empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Get the tool context
    pub fn context(&self) -> ToolContext {
        self.context.read().unwrap().clone()
    }

    /// Return a clone of the registry's artifact store handle.
    pub fn artifact_store(&self) -> ArtifactStore {
        self.artifact_store.clone()
    }

    /// Return the session-scoped immutable-content adapter, when configured.
    pub fn immutable_content_adapter(&self) -> Option<&ImmutableContentAdapterSession> {
        self.immutable_content_adapter.as_ref()
    }

    /// Get a stored tool artifact by URI.
    pub fn get_artifact(&self, artifact_uri: &str) -> Option<ToolArtifact> {
        self.artifact_store.get(artifact_uri)
    }

    /// Replace the trace sink used for compact tool/program execution events.
    pub fn set_trace_sink(&self, sink: Arc<dyn TraceSink>) {
        *self.trace_sink.write().unwrap() = sink;
    }

    /// Return the current trace sink.
    pub fn trace_sink(&self) -> Arc<dyn TraceSink> {
        Arc::clone(&self.trace_sink.read().unwrap())
    }

    /// Set the search configuration for the tool context
    pub fn set_search_config(&self, config: crate::config::SearchConfig) {
        let mut ctx = self.context.write().unwrap();
        *ctx = ctx.clone().with_search_config(config);
    }

    /// Set a sandbox executor so that `bash` tool calls use the sandbox even
    /// when executed without an explicit `ToolContext` (i.e., via `execute()`).
    pub fn set_sandbox(&self, sandbox: std::sync::Arc<dyn crate::sandbox::BashSandbox>) {
        let mut ctx = self.context.write().unwrap();
        *ctx = ctx.clone().with_sandbox(sandbox);
    }

    /// Set environment overrides used by subprocess-backed tools when executed
    /// without an explicit context.
    pub fn set_command_env(&self, env: Arc<HashMap<String, String>>) {
        let mut ctx = self.context.write().unwrap();
        *ctx = ctx.clone().with_command_env(env);
    }

    /// Execute a tool by name using the registry's default context.
    ///
    /// This is the lowest-level standalone registry boundary. It does not run
    /// agent/session permission, HITL, hook, budget, queue, timeout,
    /// cancellation, or sanitization policy.
    pub async fn execute(&self, name: &str, args: &serde_json::Value) -> Result<ToolResult> {
        let ctx = self.context();
        self.execute_with_context(name, args, &ctx).await
    }

    /// Execute a tool by name with an external caller-owned context.
    ///
    /// This remains a low-level ungoverned call; agent/session paths must use
    /// their scoped invocation gateway instead.
    pub async fn execute_with_context(
        &self,
        name: &str,
        args: &serde_json::Value,
        ctx: &ToolContext,
    ) -> Result<ToolResult> {
        let start = std::time::Instant::now();
        let prepared = self.prepare_output_with_context(name, args, ctx).await?;
        let transform_binding = prepared.transform_binding;
        let mut result = match prepared.output {
            Some(output) => Ok(ToolResult {
                name: name.to_string(),
                output: output.content,
                exit_code: if output.success { 0 } else { 1 },
                metadata: output.metadata,
                images: output.images,
                error_kind: output.error_kind,
                trust: crate::tools::ToolResultTrustV1::WorkspaceData,
            }),
            None => Ok(ToolResult::error(name, format!("Unknown tool: {name}"))),
        };

        if let Ok(result) = &mut result {
            result.metadata = Some(super::ensure_tool_result_evidence_with_transform_binding(
                result.metadata.take(),
                &result.output,
                &transform_binding,
            )?);
        }

        if let Ok(ref r) = result {
            crate::telemetry::record_tool_result(r.exit_code, start.elapsed());
            self.record_trace_event(name, r, start.elapsed());
        }

        result
    }

    /// Execute a tool and return the `ToolOutput` projection using the registry's
    /// default context. The same observation pipeline as `execute` is applied.
    pub async fn execute_raw(
        &self,
        name: &str,
        args: &serde_json::Value,
    ) -> Result<Option<ToolOutput>> {
        let ctx = self.context();
        self.execute_raw_with_context(name, args, &ctx).await
    }

    /// Execute a tool and return the `ToolOutput` projection with an external
    /// context. The same observation pipeline as `execute_with_context` is applied.
    pub async fn execute_raw_with_context(
        &self,
        name: &str,
        args: &serde_json::Value,
        ctx: &ToolContext,
    ) -> Result<Option<ToolOutput>> {
        Ok(self
            .prepare_output_with_context(name, args, ctx)
            .await?
            .output)
    }

    async fn prepare_output_with_context(
        &self,
        name: &str,
        args: &serde_json::Value,
        ctx: &ToolContext,
    ) -> Result<PreparedToolOutput> {
        let policy = self.transform_policy.read().unwrap().clone();
        let transform_binding = ToolResultTransformBindingV1::from_policy(&policy)?;
        let Some(tool) = self.get(name) else {
            return Ok(PreparedToolOutput {
                output: None,
                transform_binding,
            });
        };

        let mut output = tool.execute(args, ctx).await?;
        self.compact_change_metadata(name, &mut output.metadata, ctx)
            .await?;
        let original_content = output.content.clone();
        let truncated = transform_tool_output_with_artifact(name, &output.content, &policy);
        output.content = truncated.content;
        let loss_mode = truncated.loss_mode;
        let projected_artifact_reference = truncated.artifact.is_some();
        let artifact = truncated.artifact.or_else(|| {
            self.immutable_content_adapter
                .as_ref()
                .map(|_| super::tool_output_artifact(name, &original_content, output.content.len()))
        });
        if let Some(mut artifact) = artifact {
            let compatibility_uri = artifact.artifact_uri.clone();
            self.store_tool_artifact(
                name,
                &original_content,
                &mut artifact,
                ImmutableContentKindV1::ToolResultOriginal,
                ctx,
            )
            .await?;
            if projected_artifact_reference {
                rewrite_projected_artifact_uri(
                    &mut output.content,
                    &compatibility_uri,
                    &artifact.artifact_uri,
                )?;
            }
            output.metadata = Some(merge_tool_output_artifact_metadata(
                output.metadata,
                &artifact,
            ));
        }
        output.metadata = Some(super::attach_tool_result_evidence_with_transform_binding(
            output.metadata,
            &original_content,
            &output.content,
            loss_mode,
            &transform_binding,
        )?);
        Ok(PreparedToolOutput {
            output: Some(output),
            transform_binding,
        })
    }

    async fn store_tool_artifact(
        &self,
        tool_name: &str,
        content: &str,
        artifact: &mut ToolOutputArtifact,
        kind: ImmutableContentKindV1,
        ctx: &ToolContext,
    ) -> Result<()> {
        if let Some(adapter) = &self.immutable_content_adapter {
            let cancellation = ctx.cancellation_token();
            let retained = tokio::select! {
                biased;
                _ = cancellation.cancelled() => {
                    return Err(ImmutableContentError::Cancelled.into());
                }
                result = adapter.put(kind, TOOL_RESULT_CONTENT_MEDIA_TYPE, content.as_bytes()) => {
                    result.map_err(|error| anyhow::anyhow!(
                        "immutable content adapter '{}' rejected Tool content: {}",
                        adapter.adapter_name(),
                        error.redacted_message(),
                    ))?
                }
            };
            artifact.artifact_uri.clone_from(&retained.uri);
            artifact.content_reference = Some(retained);
            return Ok(());
        }

        self.artifact_store
            .put_content_addressed(ToolArtifact {
                artifact_id: artifact.artifact_id.clone(),
                artifact_uri: artifact.artifact_uri.clone(),
                tool_name: tool_name.to_string(),
                content: content.to_string(),
                original_bytes: artifact.original_bytes,
                shown_bytes: artifact.shown_bytes,
            })
            .map_err(|error| anyhow::anyhow!("immutable Tool artifact conflict: {error}"))?;
        Ok(())
    }

    async fn compact_change_metadata(
        &self,
        tool_name: &str,
        metadata: &mut Option<serde_json::Value>,
        ctx: &ToolContext,
    ) -> Result<()> {
        let Some(serde_json::Value::Object(object)) = metadata.as_mut() else {
            return Ok(());
        };
        let before = object
            .get("before")
            .and_then(serde_json::Value::as_str)
            .map(ToString::to_string);
        let after = object
            .get("after")
            .and_then(serde_json::Value::as_str)
            .map(ToString::to_string);
        if before.is_none() && after.is_none() {
            return Ok(());
        }

        let before_bytes = before.as_ref().map_or(0, String::len);
        let after_bytes = after.as_ref().map_or(0, String::len);
        let total_bytes = before_bytes.saturating_add(after_bytes);
        let compacted = total_bytes > MAX_INLINE_CHANGE_BYTES;
        let before_artifact = match before.as_deref() {
            Some(content) => {
                self.store_change_artifact(
                    tool_name,
                    "before",
                    content,
                    compacted,
                    ImmutableContentKindV1::ToolChangeBefore,
                    ctx,
                )
                .await?
            }
            None => None,
        };
        let after_artifact = match after.as_deref() {
            Some(content) => {
                self.store_change_artifact(
                    tool_name,
                    "after",
                    content,
                    compacted,
                    ImmutableContentKindV1::ToolChangeAfter,
                    ctx,
                )
                .await?
            }
            None => None,
        };

        let unified_diff = if compacted && total_bytes <= MAX_DIFF_COMPUTE_BYTES {
            let diff = similar::TextDiff::from_lines(
                before.as_deref().unwrap_or_default(),
                after.as_deref().unwrap_or_default(),
            )
            .unified_diff()
            .context_radius(3)
            .header("before", "after")
            .to_string();
            Some(bounded_head_tail(&diff, CHANGE_DIFF_PREVIEW_BYTES))
        } else {
            None
        };

        if compacted {
            if let Some(content) = before.as_deref() {
                object.insert(
                    "before".to_string(),
                    serde_json::Value::String(bounded_head_tail(
                        content,
                        CHANGE_SIDE_PREVIEW_BYTES,
                    )),
                );
            }
            if let Some(content) = after.as_deref() {
                object.insert(
                    "after".to_string(),
                    serde_json::Value::String(bounded_head_tail(
                        content,
                        CHANGE_SIDE_PREVIEW_BYTES,
                    )),
                );
            }
        }

        object.insert(
            "change".to_string(),
            serde_json::json!({
                "compacted": compacted,
                "before": before.as_deref().map(|content| serde_json::json!({
                    "bytes": content.len(),
                    "sha256": sha256::digest(content.as_bytes()),
                    "artifact": before_artifact,
                })),
                "after": after.as_deref().map(|content| serde_json::json!({
                    "bytes": content.len(),
                    "sha256": sha256::digest(content.as_bytes()),
                    "artifact": after_artifact,
                })),
                "unified_diff": unified_diff,
                "diff_omitted": compacted && total_bytes > MAX_DIFF_COMPUTE_BYTES,
            }),
        );
        Ok(())
    }

    async fn store_change_artifact(
        &self,
        tool_name: &str,
        side: &str,
        content: &str,
        store: bool,
        kind: ImmutableContentKindV1,
        ctx: &ToolContext,
    ) -> Result<Option<serde_json::Value>> {
        if !store
            || self.immutable_content_adapter.is_none()
                && content.len() > self.artifact_store.limits().max_bytes
        {
            return Ok(None);
        }
        let mut artifact = tool_output_artifact(&format!("{tool_name}-{side}"), content, 0);
        self.store_tool_artifact(tool_name, content, &mut artifact, kind, ctx)
            .await?;
        let mut metadata = serde_json::json!({
            "artifact_id": artifact.artifact_id,
            "artifact_uri": artifact.artifact_uri,
        });
        if let Some(reference) = artifact.content_reference {
            metadata["content_reference"] = serde_json::json!(reference);
        }
        Ok(Some(metadata))
    }

    fn record_trace_event(&self, name: &str, result: &ToolResult, duration: std::time::Duration) {
        let sink = self.trace_sink();
        sink.record(TraceEvent::tool_execution(
            name,
            result.exit_code == 0,
            result.exit_code,
            duration,
            result.output.len(),
            result.metadata.as_ref(),
        ));

        if name == "program" {
            sink.record(TraceEvent::program_execution(
                name,
                result.exit_code == 0,
                result.exit_code,
                duration,
                result.output.len(),
                result.metadata.as_ref(),
            ));
        }
    }
}

fn rewrite_projected_artifact_uri(
    projected: &mut String,
    compatibility_uri: &str,
    retained_uri: &str,
) -> Result<()> {
    if compatibility_uri == retained_uri {
        return Ok(());
    }
    let start = projected.rfind(compatibility_uri).ok_or_else(|| {
        anyhow::anyhow!(
            "Tool result projection lost its compatibility artifact reference before retention"
        )
    })?;
    let end = start + compatibility_uri.len();
    projected.replace_range(start..end, retained_uri);
    Ok(())
}

fn bounded_head_tail(content: &str, max_bytes: usize) -> String {
    if content.len() <= max_bytes {
        return content.to_string();
    }
    let head_limit = max_bytes / 2;
    let tail_limit = max_bytes.saturating_sub(head_limit);
    let head = crate::text::truncate_utf8(content, head_limit);
    let mut tail_start = content.len().saturating_sub(tail_limit);
    while tail_start < content.len() && !content.is_char_boundary(tail_start) {
        tail_start += 1;
    }
    format!(
        "{}\n\n... [{} bytes omitted from middle] ...\n\n{}",
        head,
        content
            .len()
            .saturating_sub(head.len())
            .saturating_sub(content.len().saturating_sub(tail_start)),
        &content[tail_start..]
    )
}

#[cfg(test)]
#[path = "registry/tests.rs"]
mod tests;