Designing a Unified, Type‑Safe, Extensible Tool System for AI Agents

Problem Background

An AI Agent needs to call various tools—read/write files, execute commands, search code, generate images, etc. When the number of tools reaches dozens and must support dynamic extension, unified management becomes a core challenge.

The solution proposes an architecture with four core goals:

Unified Management : Register, discover, and invoke all tools through a single mechanism.

Type Safety : Strongly typed input and output for each tool to reduce runtime errors.

Extensibility : Runtime dynamic registration of external tools without modifying core code.

Runtime Efficiency : Streamed output so long operations can return progress without waiting for completion.

Flexible Configuration : Clients can customize tool names, parameter names, enable/disable tools, and set parameters.

Design Philosophy

Type Erasure

Each tool is developed with strong input and output types, but when registered the type information is erased to JSON. The registry only sees "tool ID + JSON input + JSON output", allowing a single interface to manage all tools without specialized logic.

The key is to capture conversion functions at registration time: how to deserialize JSON to the tool's input type and how to serialize the output back to JSON.

Dependency Injection

Tools need runtime resources such as a terminal backend, file system, working directory, environment variables, and notification handles. These resources are stored in a type‑safe container rather than passed as function parameters. Tools retrieve what they need from the container at execution time, keeping the tool signature stable and allowing new resources to be added without changing existing tools.

Layered Architecture

┌──────────────────────────────────────────────────────────┐
│  Session Layer Adapter                                   │
│  Exposes the entry point, binds a session                │
│  Responsibilities: tool definition query, tool call,    │
│  dynamic register/unregister, prompt rendering, resource │
│  read/write                                            │
├──────────────────────────────────────────────────────────┤
│  Immutable Tool Set                                      │
│  Produced by the builder, immutable at runtime (except │
│  external tools)                                        │
│  Responsibilities: tool lookup, param name conversion,   │
│  context construction, post‑processing                 │
├──────────────────────────────────────────────────────────┤
│  Tool Set Builder                                        │
│  Registers all built‑in tools, accepts external config   │
│  Responsibilities: tool registration, config validation,│
│  dependency checks, produce immutable set               │
├──────────────────────────────────────────────────────────┤
│  Tool Interface + Metadata Interface                     │
│  Tool Interface: defines input/output types + exec logic│
│  Metadata Interface: defines category, namespace,          │
│  description template                                   │
├──────────────────────────────────────────────────────────┤
│  In‑process Execution Registry                           │
│  Mapping: tool ID → type‑erased execution handle        │
│  Responsibilities: find tool, execute call, return      │
│  streamed output                                       │
├──────────────────────────────────────────────────────────┤
│  Dependency‑Injection Container                          │
│  Type → value type‑safe storage                         │
│  Responsibilities: store terminal, file system,       │
│  working directory, env vars, notifications, etc.       │
└──────────────────────────────────────────────────────────┘

Tool Interface Design

4.1 Tool Interface

Associated Types

Args : Parameter structure, must be deserializable from JSON and generate a JSON Schema for the model.

Output : Return structure, must be serializable to JSON and provide a method to present the output to the model.

Required Methods id() – returns a unique identifier string, e.g., "read_file" or "bash". description(ctx) – returns name and description sent to the model.

Optional Overrides capabilities() – returns capability flags (read‑only, scope, etc.). should_list(ctx) – whether to show the tool in the current round (default true).

Execution Methods (choose one) run(ctx, args) -> Result<Output, Error> – blocking execution. execute(ctx, args) -> Stream<Output> – streaming execution for long‑running operations.

4.2 Streaming Output Model

Tool execution returns an asynchronous stream: [Progress, Progress, ..., Terminal] Progress : Zero or more intermediate state updates (e.g., terminal output fragments) for real‑time progress display.

Terminal : Exactly one final item containing the result (success or error).

4.3 Type‑Erasure Adapter

The system provides a type‑erased version of the tool interface. The compiler generates conversion functions that:

Deserialize JSON parameters to the concrete input type.

Invoke the tool's execution method.

Serialize the output back to JSON.

Extract the "content for the model" and optional chat‑completion response.

Wrap everything into a unified type‑erased output.

4.4 Metadata Interface

kind()

– high‑level category (e.g., FileOperation, CommandExecution). namespace() – logical grouping; the same functionality can have different implementations in different namespaces. description_template() – template string for description, supports placeholder substitution.

Optional overrides: is_read_only(), emitted_notifications(), requires_expr(), versioned_definition(...).

Tool Classification System

File Operations : Read, Edit, Write, Delete, Move, ListDir.

Code Understanding : Search, Lsp.

Command Execution : Execute.

Network : WebSearch, WebFetch.

Task Management : BackgroundTaskAction, WaitTasksAction, KillTaskAction, Task.

Media Generation : ImageGen, VideoGen, ImageToVideo.

Interaction : AskUser, Plan.

Other : MemorySearch, Skill, Monitor, Other.

The three core uses of classification are template rendering, permission inference (read‑only detection), and capability filtering.

Dependency‑Injection Container

The container is a type‑safe heterogeneous map (type → value). It provides: insert<T>(value) – store a value of type T. get<T>() -> Option<T> – retrieve by type. require<T>() -> T – retrieve or error. get_or_default<T>() -> &mut T – retrieve or insert a default.

The container is thread‑safe (protected by a mutex) and holds resources such as terminal backend, file system, working directory, environment variables, notification handles, web‑search client, language‑service backend, media‑generation client, tool configuration, and runtime state.

Configuration vs. Runtime State

Params<T> : Configuration parameters injected at tool‑set initialization and immutable during runtime (e.g., max timeout, background‑execution flag).

State<T> : Runtime state that can change during execution and is persisted to disk (e.g., todo list, task completion records).

Tool Registration

7.1 Builder

The tool‑set builder registers all built‑in tools at initialization. Tools are registered either without configuration parameters or with a configuration type P.

7.2 Registration Process

Create a tool instance via its default constructor.

Obtain namespace, tool ID, and category from the metadata interface.

Compose the fully‑qualified ID "namespace:toolID" (e.g., "Main:bash").

Generate JSON Schema from the input type for the model.

Create conversion functions (output converter, parameter validator, parameter applier, input parser, local registrar).

Package all information into a tool entry and store it in the registry.

7.3 External Tool Packages

External tool packages can be registered at process start via a callback that receives the builder reference, allowing addition without core code changes.

7.4 Reminder Registration

Reminders are cross‑tool additional logic evaluated after each tool call. They have dependency expressions and activate only when conditions are met.

Configuration and Validation

8.1 Tool‑Set Configuration

Clients specify which tools to enable and how to configure them. Each entry includes tool ID, JSON parameters, optional name override, parameter name remapping, description override, behavior version, and optional classification.

8.2 Validation Process

Tool existence: every configured tool ID must exist in the registry.

Behavior version: specified version must be registered.

Parameter validation: client‑provided parameters must pass the tool's validator.

Name uniqueness: client‑visible names must be unique.

Tool‑set conflicts: mutually exclusive variants cannot coexist.

Dependency checks: evaluate dependency expressions between tools.

8.3 Dependency Checks

Expressions can reference namespace, tool ID, category, parameter values, or input schema. Example rules: an edit tool may require a read tool, plan‑mode tools require a corresponding exit‑plan tool, background‑task output tools require a tool that can produce background tasks.

Initialization Flow

Validate configuration (steps above).

Build a template renderer that maps category → client tool name for placeholder substitution.

Create the dependency‑injection container and inject session resources (terminal, file system, working directory, etc.).

Register each tool’s configuration and runtime‑state serialization information.

Load persisted runtime state from disk.

For each configured tool:

Extract the tool entry from the builder.

Register the tool instance in the in‑process execution registry.

Resolve behavior version.

Merge default parameters with client overrides (client overrides win).

Render the tool definition sent to the model (description template + parameter name remapping).

Write effective parameters into the DI container.

Package the final tool entry.

Filter active reminders based on dependency expressions.

Start background scheduler if present.

Return the immutable tool set.

Tool Call Flow

10.1 End‑to‑End Call Path

Large model outputs tool_call (tool name + args + call ID)
    │
    ▼
Session Layer Adapter receives the call
    │
    ▼
Immutable Tool Set: prepare dispatch
    │  1. Find tool by client name in the tool list
    │  2. If parameter name remapping exists, translate client names back to canonical names
    │  3. Build call context (inject DI container, template renderer, work dir, etc.)
    │  4. Look up the tool's execution handle in the in‑process registry
    │
    ▼
Execution Handle: type‑erased execution
    │  1. Deserialize JSON args to the concrete input type
    │  2. Call the tool's execution method
    │  3. Tool fetches the DI container from the context
    │  4. Read required backends (terminal, file system, …)
    │  5. Perform the actual operation
    │  6. Return tool output
    │  7. Serialize output to JSON
    │  8. Extract "content for the model"
    │
    ▼
Immutable Tool Set: post‑processing
    │  1. Convert JSON output back to the unified output enum
    │  2. Evaluate reminders (if enabled)
    │  3. Generate final text (tool output + reminder text)
    │  4. Persist runtime state to disk
    │
    ▼
Return execution result (structured output + model text + metadata)
    │
    ▼
Session Layer passes result back to the large model

10.2 Streaming Calls

Blocking mode : wait for completion, collect all progress items, then return the terminal item.

Streaming mode : immediately return an async stream; the caller receives progress items in real time and the final terminal item when done.

10.3 Meta‑Tools and Internal Dispatch

Meta‑tools can invoke other tools. To avoid double post‑processing, meta‑tools use an internal dispatch channel that skips the reminder and state‑persistence steps, letting the outer call handle them once.

Dynamic Tool Registration

11.1 External Tools

The immutable tool set supports runtime registration of external tools (e.g., discovered via MCP protocol). External tools have JSON‑only input/output, belong to the "Other" category, use a separate namespace, and obtain their input schema from a remote server.

11.2 Naming Conventions

External tool client names often carry a special prefix to distinguish them from built‑in tools; queries can filter them out.

11.3 Configuration Merge Pipeline

Obtain baseline configuration from session defaults.

Merge external tool snapshots (skip ID conflicts).

Merge remote Hub tools (skip conflicts with baseline or external).

Capability filtering based on mode (e.g., read‑only only).

Initialize the tool set with the merged configuration.

Priority: baseline > external > remote.

Reminder System

12.1 Design

Reminders are additional messages appended after tool execution to guide the model (e.g., "you have an unfinished todo", "code diagnostics found, consider checking", "new skill available").

12.2 Types

Cross‑tool reminders : registered in the builder, evaluated after every tool call, have their own dependency expression.

Tool‑internal reminders : declared in a tool’s metadata and evaluated by the tool itself.

12.3 Execution

During post‑processing, if reminders are enabled, the system iterates over active reminders, calls their collect method with the DI container and tool output, and appends the returned text fragments (wrapped with special tags) to the tool output.

Session Layer Adapter

Tool definition query : returns definitions (name, description, JSON schema) of all enabled tools for the model.

Tool call : receives a tool call from the model, executes it, and returns the result.

Dynamic register/unregister : add or remove external tools at runtime.

Prompt rendering : render prompt templates with placeholders.

Resource read/write : access resources stored in the DI container.

A key detail: the terminal backend is stored separately from the tool‑set lock, allowing immediate termination of a running process when the user cancels.

Error Handling

14.1 Tool Execution Errors

Errors contain a type tag and a human‑readable message, which are returned to the model for context.

Types include: unimplemented, invalid parameters, not found, permission denied, unauthorized, timeout, cancelled, rate‑limited, execution error, custom error.

14.2 Configuration Validation Errors

Validation errors include tool ID, error message, field path, expected value, and actual value, with a classification tag for filtering.

Tool Implementation Example (Command Execution Tool)

Input Definition

Command Execution Input:
command: String          — command to execute
timeout: Number?        — optional timeout in ms
description: String     — command description
is_background: Bool    — run in background?

Output Definition

Command Execution Output (two possibilities):
  Foreground result: { output_text, exit_code, truncated?, ... }
  Background started: { task_id }

Tool Interface Implementation

id()

returns "bash". capabilities() returns { is_read_only: false, scope: Write }. run(ctx, input) performs:

Obtain DI container from context.

Read terminal backend, working directory, env vars, notification handle.

Build execution request (command, cwd, timeout, output limits).

Invoke terminal backend to run the command.

Send completion notification.

Return the result.

Metadata Implementation

kind()

returns Execute. namespace() returns the main tool‑set namespace. description_template() returns a template string. emitted_notifications() returns ["execution_complete", "execution_failed", "execution_timeout", ...].

Configuration Parameters

Command Execution Config:
  enabled_background: Bool   — allow background execution
  max_timeout_ms: Number    — maximum timeout limit

These parameters are written into the DI container at initialization and read by the tool during execution (e.g., to clamp user‑provided timeout).

Core Data Structures

Tool Definition (sent to model)

Tool Definition:
type: "function"
function:
    name: String            — client‑visible name
    description: String?    — tool description
    parameters: JSONSchema   — parameter schema

Unified Input Enum

All tool inputs are represented by a single enum; each variant holds the concrete input struct, plus a "dynamic" variant for external tools that simply passes JSON.

Unified Output Enum

All tool outputs are represented by a single enum; each variant holds the concrete output struct, plus "dynamic" and "plain_text" variants for generic cases.

Execution Result

Execution Result:
  output: UnifiedOutputEnum          — structured output
  prompt_text: String                — text sent to the model (including reminders)
  effective_tool_name: String?      — actual target tool name when a meta‑tool is used

Key Design Decisions

Two‑layer interface (execution + metadata) separates runtime logic from descriptive data, avoiding pollution.

Type erasure occurs at registration; conversion functions are captured once, allowing the dispatcher to work purely with JSON.

Dependency‑injection container prevents signature bloat; tools always have the same signature (context, input) and pull needed resources from the container.

Builder pattern produces an immutable tool set; only external tools can be added at runtime.

Execution registry holds tool instances, while the tool set manages metadata and configuration, keeping lookup responsibilities distinct.

Streaming output enables long‑running operations to provide progress without blocking.

Bidirectional parameter name remapping lets clients use custom names while the tool internally uses canonical names.

Tool‑dependency validation via expressions ensures a self‑consistent tool set before startup.

External tools are dynamically registered, use JSON pass‑through, belong to the "Other" category, and have separate namespaces.

Runtime state persistence automatically saves tool state to disk and restores it on next start.

Behavior versioning allows tools to evolve without breaking existing clients.

Terminal backend is stored outside the tool‑set lock, enabling immediate cancellation of running processes.

References

xai‑org/grok‑build: https://github.com/xai-org/grok-build agiwo: https://github.com/xhwSkhizein/agiwo