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,
}
pub struct ToolRegistry {
tools: RwLock<HashMap<String, Arc<dyn Tool>>>,
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>,
}
impl ToolRegistry {
pub fn new(workspace: PathBuf) -> Self {
Self::with_artifact_limits(workspace, ArtifactStoreLimits::default())
}
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),
)
}
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()),
}
}
pub(crate) fn snapshot_with_external_tools(
&self,
external: impl IntoIterator<Item = Arc<dyn Tool>>,
) -> Result<Self, ToolRegistrySnapshotError> {
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(())
}
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);
}
pub fn register(&self, tool: 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;
}
tracing::debug!("Registering tool: {}", name);
tools.insert(name, tool);
}
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))
}
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
}
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
}
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()
}
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);
}
pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
let tools = self.tools.read().unwrap();
tools.get(name).cloned()
}
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))
}
pub fn contains(&self, name: &str) -> bool {
let tools = self.tools.read().unwrap();
tools.contains_key(name)
}
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
}
pub fn list(&self) -> Vec<String> {
let tools = self.tools.read().unwrap();
let mut names = tools.keys().cloned().collect::<Vec<_>>();
names.sort();
names
}
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())
}
pub fn len(&self) -> usize {
let tools = self.tools.read().unwrap();
tools.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn context(&self) -> ToolContext {
self.context.read().unwrap().clone()
}
pub fn artifact_store(&self) -> ArtifactStore {
self.artifact_store.clone()
}
pub fn immutable_content_adapter(&self) -> Option<&ImmutableContentAdapterSession> {
self.immutable_content_adapter.as_ref()
}
pub fn get_artifact(&self, artifact_uri: &str) -> Option<ToolArtifact> {
self.artifact_store.get(artifact_uri)
}
pub fn set_trace_sink(&self, sink: Arc<dyn TraceSink>) {
*self.trace_sink.write().unwrap() = sink;
}
pub fn trace_sink(&self) -> Arc<dyn TraceSink> {
Arc::clone(&self.trace_sink.read().unwrap())
}
pub fn set_search_config(&self, config: crate::config::SearchConfig) {
let mut ctx = self.context.write().unwrap();
*ctx = ctx.clone().with_search_config(config);
}
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);
}
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);
}
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
}
pub async fn execute_with_context(
&self,
name: &str,
args: &serde_json::Value,
ctx: &ToolContext,
) -> Result<ToolResult> {
let start = std::time::Instant::now();
let policy = self.transform_policy.read().unwrap().clone();
let transform_binding = ToolResultTransformBindingV1::from_policy(&policy)?;
let tool = self.get(name);
let mut result = match tool {
Some(tool) => {
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(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,
})
}
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
}
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
}
pub async fn execute_raw_with_context(
&self,
name: &str,
args: &serde_json::Value,
ctx: &ToolContext,
) -> Result<Option<ToolOutput>> {
let policy = self.transform_policy.read().unwrap().clone();
let transform_binding = ToolResultTransformBindingV1::from_policy(&policy)?;
let tool = self.get(name);
match tool {
Some(tool) => {
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(Some(output))
}
None => Ok(None),
}
}
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(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,
});
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;