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