Custom Code Tool - Developer Toolkit
This sample code demonstrates how you can use aidputils to test your Custom Code tool.
The Developer Toolkit example demonstrates a multi-tool package and the use of helper modules in a utils/ directory. The package registers three tools — a bash command runner, a file operations tool, and a Python code runner — and uses shared helper functions for output truncation and path sanitization.
Note:
The Developer Toolkit is an illustrative example. Bash command execution and Python code execution have significant security implications. In production, restrict the AI compute, sandbox the operations, and apply strict allow-lists for the commands and code patterns the tool will execute.Package Layout
advanced_tool.zip
├── tool_implementation.py
├── tool_config.json
├── requirements.txt # stdlib only
└── utils/
├── __init__.py
└── text_utils.py # truncate_output, sanitize_path
tool_implementation.py
import subprocess
import os
from aidputils.agents.tools.custom_tools.base import CustomToolBase
from .utils.text_utils import truncate_output, sanitize_path
def _get_cfg(conf, key, default):
"""Read a config value from either the outer dict or the
nested user conf. Coerces numeric settings to int to avoid
type mismatches when values are rendered as strings by the
template substitution layer."""
inner = conf.get("conf") if isinstance(conf, dict) else None
if isinstance(inner, dict) and key in inner:
value = inner[key]
elif isinstance(conf, dict) and key in conf:
value = conf[key]
else:
value = default
if isinstance(default, int) and not isinstance(value, bool):
try:
return int(value)
except (TypeError, ValueError):
return default
return value
@BaseTool.register
class BashTool(CustomToolBase):
"""Execute bash commands and return output."""
@classmethod
def _execute_tool(cls, conf, runtime_params, **context_vars):
command = runtime_params.get("command", "")
timeout = _get_cfg(conf, "timeout", 30)
max_lines = _get_cfg(conf, "max_output_lines", 200)
try:
result = subprocess.run(
["bash", "-c", command],
capture_output=True, text=True, timeout=timeout
)
except subprocess.TimeoutExpired:
# Surface the timeout as a tool failure rather than
# returning {"error": ...}, which would be treated as
# a successful response.
raise RuntimeError(f"Command timed out after {timeout}s")
output = result.stdout or ""
if result.stderr:
output += "\n[stderr]\n" + result.stderr
return {"output": truncate_output(output, max_lines)}
@BaseTool.register
class FileTool(CustomToolBase):
"""Read, write, or list files in the workspace."""
@classmethod
def _execute_tool(cls, conf, runtime_params, **context_vars):
operation = runtime_params.get("operation", "")
path = runtime_params.get("path", "")
content = runtime_params.get("content", "")
base_dir = _get_cfg(conf, "base_dir", "/workspace")
max_size = _get_cfg(conf, "max_file_size_kb", 1024) * 1024
safe_path = sanitize_path(base_dir, path)
if safe_path is None:
raise ValueError("Invalid path: path traversal detected")
if operation == "read":
with open(safe_path, "r") as f:
return {"output": f.read()}
if operation == "write":
parent = os.path.dirname(safe_path)
if parent:
os.makedirs(parent, exist_ok=True)
with open(safe_path, "w") as f:
f.write(content)
return {"output": f"Written {len(content)} chars to {path}"}
if operation == "list":
target = safe_path if os.path.isdir(safe_path) else os.path.dirname(safe_path)
return {"output": "\n".join(sorted(os.listdir(target)))}
raise ValueError(f"Unknown operation: {operation}. Use read/write/list")
@BaseTool.register
class PythonTool(CustomToolBase):
"""Execute Python code in an isolated subprocess."""
@classmethod
def _execute_tool(cls, conf, runtime_params, **context_vars):
code = runtime_params.get("code", "")
timeout = _get_cfg(conf, "timeout", 60)
max_lines = _get_cfg(conf, "max_output_lines", 500)
try:
result = subprocess.run(
["python3", "-c", code],
capture_output=True, text=True, timeout=timeout
)
except subprocess.TimeoutExpired:
raise RuntimeError(f"Execution timed out after {timeout}s")
output = result.stdout or ""
if result.stderr:
output += "\n[stderr]\n" + result.stderr
return {"output": truncate_output(output, max_lines)}
tool_config.json
{
"displayName": "Developer Toolkit",
"description": "A collection of tools for bash commands, file operations, and Python execution",
"tools": [
{
"toolClassName": "BashTool",
"displayName": "Bash Tool",
"description": "Executes a bash command and returns stdout/stderr output",
"version": "1.0.0",
"schema": [
{
"name": "command",
"type": "string",
"description": "The bash command to execute"
}
],
"conf": {
"timeout": 30,
"max_output_lines": 200
}
},
{
"toolClassName": "FileTool",
"displayName": "File Tool",
"description": "Read, write, or list files in the workspace",
"version": "1.0.0",
"schema": [
{"name": "operation", "type": "string",
"description": "Operation to perform: read, write, or list"},
{"name": "path", "type": "string",
"description": "File or directory path"},
{"name": "content", "type": "string",
"description": "Content to write (for write operation)"}
],
"conf": {
"base_dir": "/workspace",
"max_file_size_kb": 1024
}
},
{
"toolClassName": "PythonTool",
"displayName": "Python Tool",
"description": "Executes Python code in an isolated subprocess and returns the output",
"version": "1.0.0",
"schema": [
{"name": "code", "type": "string",
"description": "The Python code to execute"}
],
"conf": {
"timeout": 60,
"max_output_lines": 500
}
}
]
}utils/text_utils.py
def truncate_output(text, max_lines=200):
if not text:
return ""
try:
max_lines = int(max_lines)
except (TypeError, ValueError):
max_lines = 200
lines = text.strip().split("\n")
if len(lines) > max_lines:
lines = lines[:max_lines] + [f"... ({len(lines) - max_lines} lines truncated)"]
return "\n".join(lines)
def sanitize_path(base_dir, relative_path):
import os
if not relative_path:
return base_dir
full = os.path.normpath(os.path.join(base_dir, relative_path))
if not full.startswith(os.path.normpath(base_dir)):
return None
return full
utils/__init__.py
# Empty file. Required for Python to treat utils/ as a package.requirements.txt
# stdlib onlyAfter uploading the ZIP, the Package tab shows the three discovered tools and lets you enable or disable each one. The Parameters tab shows a Tool Class drop-down that switches between BashTool, FileTool, and PythonTool, and exposes the per-tool configuration (timeout, max_output_lines, base_dir, max_file_size_kb) on the right.