showcase/shell-docs/src/content/ag-ui/sdk/kotlin/tools/tool-registry.mdx
ToolRegistry manages tool executors and handles tool execution with timeout handling, statistics tracking, and thread-safe access. It serves as the central registry for all available client-side tools.
interface ToolRegistry {
fun registerTool(executor: ToolExecutor)
fun unregisterTool(toolName: String): Boolean
fun getToolExecutor(toolName: String): ToolExecutor?
fun getAllTools(): List<Tool>
fun getAllExecutors(): Map<String, ToolExecutor>
fun isToolRegistered(toolName: String): Boolean
suspend fun executeTool(context: ToolExecutionContext): ToolExecutionResult
fun getToolStats(toolName: String): ToolExecutionStats?
fun getAllStats(): Map<String, ToolExecutionStats>
fun clearStats()
}
val toolRegistry = toolRegistry {
addTool(CalculatorToolExecutor())
addTool(WeatherToolExecutor())
addTool(FileToolExecutor())
}
val toolRegistry = toolRegistry(
CalculatorToolExecutor(),
WeatherToolExecutor(),
FileToolExecutor()
)
val registry: ToolRegistry = DefaultToolRegistry()
registry.registerTool(CalculatorToolExecutor())
registry.registerTool(WeatherToolExecutor())
val calculator = CalculatorToolExecutor()
toolRegistry.registerTool(calculator)
// Throws IllegalArgumentException if tool name already exists
try {
toolRegistry.registerTool(AnotherCalculatorExecutor())
} catch (e: IllegalArgumentException) {
println("Tool already registered: ${e.message}")
}
val wasRemoved = toolRegistry.unregisterTool("calculator")
if (wasRemoved) {
println("Calculator tool removed")
} else {
println("Calculator tool was not registered")
}
// Check if tool exists
if (toolRegistry.isToolRegistered("weather")) {
println("Weather tool is available")
}
// Get specific executor
val weatherExecutor = toolRegistry.getToolExecutor("weather")
weatherExecutor?.let { executor ->
println("Found executor: ${executor.tool.description}")
}
// Get all tools for agent registration
val allTools = toolRegistry.getAllTools()
println("Available tools: ${allTools.map { it.name }}")
The registry handles tool execution automatically when used with agents, but can also be called directly:
val toolCall = ToolCall(
id = "calc-1",
function = ToolCall.Function(
name = "calculator",
arguments = buildJsonObject {
put("expression", "2 + 3")
}
)
)
val context = ToolExecutionContext(
toolCall = toolCall,
threadId = "thread-123",
runId = "run-456"
)
try {
val result = toolRegistry.executeTool(context)
if (result.success) {
println("Result: ${result.result}")
println("Message: ${result.message}")
} else {
println("Execution failed: ${result.message}")
}
} catch (e: ToolNotFoundException) {
println("Tool not found: ${e.message}")
} catch (e: ToolExecutionException) {
println("Execution error: ${e.message}")
}
data class ToolExecutionStats(
val executionCount: Long = 0,
val successCount: Long = 0,
val failureCount: Long = 0,
val totalExecutionTimeMs: Long = 0,
val averageExecutionTimeMs: Double = 0.0
) {
val successRate: Double get() = if (executionCount > 0) successCount.toDouble() / executionCount else 0.0
}
val calculatorStats = toolRegistry.getToolStats("calculator")
calculatorStats?.let { stats ->
println("Calculator executions: ${stats.executionCount}")
println("Success rate: ${String.format("%.1f%%", stats.successRate * 100)}")
println("Average execution time: ${stats.averageExecutionTimeMs}ms")
println("Total execution time: ${stats.totalExecutionTimeMs}ms")
}
val allStats = toolRegistry.getAllStats()
allStats.forEach { (toolName, stats) ->
println("$toolName: ${stats.executionCount} executions, ${String.format("%.1f%%", stats.successRate * 100)} success rate")
}
// Clear all statistics
toolRegistry.clearStats()
println("All statistics cleared")
try {
val context = ToolExecutionContext(unknownToolCall)
toolRegistry.executeTool(context)
} catch (e: ToolNotFoundException) {
println("Tool '${e.message}' is not registered")
// Could suggest similar tool names or prompt for registration
}
try {
val result = toolRegistry.executeTool(context)
if (!result.success) {
println("Tool execution failed: ${result.message}")
// Handle graceful failure
}
} catch (e: ToolExecutionException) {
println("Unrecoverable error: ${e.message}")
println("Tool: ${e.toolName}, Call ID: ${e.toolCallId}")
// Log error for debugging
}
Timeouts are handled automatically based on each tool's getMaxExecutionTimeMs():
class SlowToolExecutor : ToolExecutor {
override fun getMaxExecutionTimeMs(): Long = 60_000 // 1 minute
override suspend fun execute(context: ToolExecutionContext): ToolExecutionResult {
// Long-running operation
delay(30_000) // 30 seconds
return ToolExecutionResult.success(JsonPrimitive("Done"))
}
}
val agent = agentWithTools(
url = "https://api.example.com/agent",
toolRegistry = toolRegistry
) {
bearerToken = "your-token"
}
// Tools are automatically available to the agent
agent.sendMessage("Calculate 2 + 3").collect { event ->
// Agent will use calculator tool automatically
}
The registry provides tool definitions to agents:
// This happens automatically in agentWithTools
val toolDefinitions = toolRegistry.getAllTools()
val input = RunAgentInput(
threadId = "thread-1",
runId = "run-1",
messages = messages,
tools = toolDefinitions // Registry tools sent to agent
)
class AgentService {
private val toolRegistry = toolRegistry {
addTool(CalculatorToolExecutor())
addTool(WeatherToolExecutor())
}
private val agent = agentWithTools(url, toolRegistry) {
bearerToken = token
}
fun getExecutionStats(): Map<String, ToolExecutionStats> {
return toolRegistry.getAllStats()
}
fun resetStats() {
toolRegistry.clearStats()
}
}
class DynamicToolManager {
private val registry: ToolRegistry = DefaultToolRegistry()
fun enableFeature(feature: String) {
when (feature) {
"location" -> registry.registerTool(LocationToolExecutor())
"camera" -> registry.registerTool(CameraToolExecutor())
"files" -> registry.registerTool(FileToolExecutor())
}
}
fun disableFeature(feature: String) {
registry.unregisterTool(feature)
}
fun getAvailableFeatures(): List<String> {
return registry.getAllTools().map { it.name }
}
}
class ToolMonitor {
private val errorCounts = mutableMapOf<String, Int>()
suspend fun executeWithMonitoring(
registry: ToolRegistry,
context: ToolExecutionContext
): ToolExecutionResult {
return try {
val result = registry.executeTool(context)
if (!result.success) {
recordError(context.toolCall.function.name)
}
result
} catch (e: ToolExecutionException) {
recordError(context.toolCall.function.name)
throw e
}
}
private fun recordError(toolName: String) {
errorCounts[toolName] = errorCounts.getOrDefault(toolName, 0) + 1
// Disable problematic tools
if (errorCounts[toolName]!! > 5) {
println("Tool $toolName has too many errors, consider disabling")
}
}
}
DefaultToolRegistry is thread-safe and can be used concurrently:
val registry = DefaultToolRegistry()
// Safe to register from multiple threads
launch { registry.registerTool(Tool1Executor()) }
launch { registry.registerTool(Tool2Executor()) }
// Safe to execute concurrently
launch { registry.executeTool(context1) }
launch { registry.executeTool(context2) }
class OptimizedToolRegistry {
private val registry = DefaultToolRegistry()
private val frequentTools = setOf("calculator", "weather", "location")
init {
// Pre-register frequently used tools
frequentTools.forEach { toolName ->
when (toolName) {
"calculator" -> registry.registerTool(CalculatorToolExecutor())
"weather" -> registry.registerTool(WeatherToolExecutor())
"location" -> registry.registerTool(LocationToolExecutor())
}
}
}
fun getOptimizedStats(): Map<String, Double> {
return registry.getAllStats().mapValues { (_, stats) ->
stats.averageExecutionTimeMs
}.toList().sortedBy { it.second }.toMap()
}
}