447 lines
15 KiB
JavaScript
447 lines
15 KiB
JavaScript
const crypto = require("crypto");
|
|
const { roleAllows } = require("./permissions");
|
|
|
|
class ToolRegistry {
|
|
constructor(audit) {
|
|
this.tools = new Map();
|
|
this.confirmations = new Map();
|
|
this.audit = typeof audit === "function" ? audit : () => {};
|
|
}
|
|
|
|
register(definition) {
|
|
validateDefinition(definition);
|
|
if (this.tools.has(definition.tool_id)) {
|
|
throw new Error(`AI tool ${definition.tool_id} is already registered.`);
|
|
}
|
|
const normalized = {
|
|
required_role: "user",
|
|
confirmation_required: true,
|
|
risk_level: "sensitive",
|
|
read_only: false,
|
|
use_cases: [],
|
|
output_expectations: "",
|
|
...definition
|
|
};
|
|
this.tools.set(normalized.tool_id, normalized);
|
|
return () => this.unregister(normalized.tool_id, normalized.owning_plugin);
|
|
}
|
|
|
|
unregister(toolId, owner = null) {
|
|
const definition = this.tools.get(toolId);
|
|
if (!definition || (owner && definition.owning_plugin !== owner)) return false;
|
|
this.tools.delete(toolId);
|
|
for (const [id, pending] of this.confirmations) {
|
|
if (pending.def?.tool_id === toolId) this.confirmations.delete(id);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
unregisterOwner(owner) {
|
|
let removed = 0;
|
|
for (const [id, definition] of this.tools) {
|
|
if (definition.owning_plugin === owner && this.unregister(id, owner)) removed += 1;
|
|
}
|
|
return removed;
|
|
}
|
|
|
|
has(toolId) {
|
|
return this.tools.has(toolId);
|
|
}
|
|
|
|
list(role) {
|
|
return [...this.tools.values()]
|
|
.filter((tool) => roleAllows(role, tool.required_role))
|
|
.map(publicDefinition);
|
|
}
|
|
|
|
inspect({ role, user, context = null } = {}) {
|
|
const decisions = [...this.tools.values()].map((definition) =>
|
|
exposureDecision(definition, { role, user, context })
|
|
);
|
|
return {
|
|
considered: decisions,
|
|
exposed: decisions.filter((decision) => decision.exposed).map((decision) => decision.tool)
|
|
};
|
|
}
|
|
|
|
validate(toolId, args, role) {
|
|
const definition = this.tools.get(toolId);
|
|
if (!definition) throw toolError("not_registered", "Tool is not registered.");
|
|
if (!roleAllows(role, definition.required_role)) {
|
|
throw toolError("permission_blocked", "Permission denied for this tool.");
|
|
}
|
|
const schema = definition.schema || {};
|
|
const input = args && typeof args === "object" && !Array.isArray(args) ? args : {};
|
|
const unknown = Object.keys(input).filter((key) => !Object.hasOwn(schema, key));
|
|
if (unknown.length) throw toolError("schema_invalid", `Unknown tool argument: ${unknown[0]}.`);
|
|
const clean = {};
|
|
for (const [key, spec] of Object.entries(schema)) {
|
|
const descriptor = typeof spec === "string" ? { type: spec, required: true } : spec || {};
|
|
const value = input[key];
|
|
if (value == null) {
|
|
if (descriptor.required === false) continue;
|
|
throw toolError("schema_invalid", `${key} is required.`);
|
|
}
|
|
clean[key] = normalizeArgument(key, value, descriptor);
|
|
}
|
|
return { def: definition, args: clean };
|
|
}
|
|
|
|
prepare({ tool, args, user, role, sessionId, context = null }) {
|
|
const definition = this.tools.get(tool);
|
|
if (!definition) {
|
|
this.recordDecision({
|
|
user,
|
|
role,
|
|
context,
|
|
tool,
|
|
args,
|
|
decision: "rejected",
|
|
rejected_reason: "not_registered"
|
|
});
|
|
throw toolError("not_registered", "Tool is not registered.");
|
|
}
|
|
const exposure = exposureDecision(definition, { role, user, context });
|
|
if (!exposure.exposed) {
|
|
this.recordDecision({
|
|
user,
|
|
role,
|
|
context,
|
|
tool,
|
|
args,
|
|
decision: "rejected",
|
|
rejected_reason: exposure.reason
|
|
});
|
|
throw toolError(exposure.reason, exposure.message);
|
|
}
|
|
let checked;
|
|
try {
|
|
checked = this.validate(tool, args, role);
|
|
} catch (error) {
|
|
this.recordDecision({
|
|
user,
|
|
role,
|
|
context,
|
|
tool,
|
|
args,
|
|
decision: "rejected",
|
|
rejected_reason: error.code || "schema_invalid"
|
|
});
|
|
throw error;
|
|
}
|
|
let allowed;
|
|
try {
|
|
allowed = checked.def.permission_check({
|
|
user,
|
|
arguments: checked.args,
|
|
required_permission: checked.def.required_permission,
|
|
context,
|
|
phase: "execute"
|
|
});
|
|
} catch {
|
|
allowed = false;
|
|
}
|
|
if (allowed && typeof allowed.then === "function") {
|
|
this.recordDecision({
|
|
user,
|
|
role,
|
|
context,
|
|
tool,
|
|
args: checked.args,
|
|
decision: "rejected",
|
|
rejected_reason: "permission_blocked"
|
|
});
|
|
throw toolError("permission_blocked", "AI tool permission checks must be synchronous.");
|
|
}
|
|
if (!allowed) {
|
|
this.recordDecision({
|
|
user,
|
|
role,
|
|
context,
|
|
tool,
|
|
args: checked.args,
|
|
decision: "rejected",
|
|
rejected_reason: "permission_blocked"
|
|
});
|
|
throw toolError("permission_blocked", "The requesting user does not have permission for this action.");
|
|
}
|
|
this.recordDecision({
|
|
user,
|
|
role,
|
|
context,
|
|
tool,
|
|
args: checked.args,
|
|
decision: "selected",
|
|
rejected_reason: null
|
|
});
|
|
if (!checked.def.confirmation_required) return { execute: true, checked, context };
|
|
const id = crypto.randomUUID();
|
|
this.confirmations.set(id, {
|
|
id,
|
|
userId: user.id,
|
|
sessionId,
|
|
expiresAt: Date.now() + 120000,
|
|
context,
|
|
...checked
|
|
});
|
|
return {
|
|
execute: false,
|
|
checked,
|
|
context,
|
|
confirmation: {
|
|
id,
|
|
display_name: checked.def.display_name,
|
|
arguments: checked.args,
|
|
expires_at: Date.now() + 120000
|
|
}
|
|
};
|
|
}
|
|
|
|
async execute({ checked, user, requestId, context = null }) {
|
|
const started = Date.now();
|
|
try {
|
|
const result = await checked.def.workflow_handler({
|
|
arguments: checked.args,
|
|
user,
|
|
ctx: context,
|
|
initiated_via_ai: true,
|
|
ai_request_id: requestId
|
|
});
|
|
this.audit({
|
|
kind: "tool",
|
|
status: "success",
|
|
user_id: user.id,
|
|
actor: auditActor(user),
|
|
origin: contextOrigin(context),
|
|
tool_requested: checked.def.tool_id,
|
|
selected_tool: checked.def.tool_id,
|
|
tool_executed: true,
|
|
arguments_summary: summarizeArguments(checked.args),
|
|
execution_ms: Date.now() - started
|
|
});
|
|
return result;
|
|
} catch (error) {
|
|
this.audit({
|
|
kind: "tool",
|
|
status: "failed",
|
|
user_id: user.id,
|
|
actor: auditActor(user),
|
|
origin: contextOrigin(context),
|
|
tool_requested: checked.def.tool_id,
|
|
selected_tool: checked.def.tool_id,
|
|
tool_executed: false,
|
|
arguments_summary: summarizeArguments(checked.args),
|
|
execution_ms: Date.now() - started,
|
|
rejected_reason: error.code || "execution_failed"
|
|
});
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
async confirm({ id, user, sessionId }) {
|
|
const pending = this.confirmations.get(id);
|
|
this.confirmations.delete(id);
|
|
if (!pending || pending.expiresAt < Date.now() || pending.userId !== user.id || pending.sessionId !== sessionId) {
|
|
throw new Error("Confirmation is invalid or expired.");
|
|
}
|
|
return this.execute({
|
|
checked: { def: pending.def, args: pending.args },
|
|
user,
|
|
requestId: id,
|
|
context: pending.context
|
|
});
|
|
}
|
|
|
|
cancel(id, userId) {
|
|
const pending = this.confirmations.get(id);
|
|
if (pending?.userId !== userId) return false;
|
|
this.confirmations.delete(id);
|
|
return true;
|
|
}
|
|
|
|
recordDecision({ user, role, context, tool, args, decision, rejected_reason }) {
|
|
this.audit({
|
|
kind: "tool_decision",
|
|
status: decision === "rejected" ? "rejected" : "selected",
|
|
user_id: user?.id || null,
|
|
actor: auditActor(user, role),
|
|
origin: contextOrigin(context),
|
|
selected_tool: decision === "selected" ? tool : null,
|
|
tool_requested: tool || null,
|
|
arguments_summary: summarizeArguments(args),
|
|
rejected_reason: rejected_reason || null
|
|
});
|
|
}
|
|
}
|
|
|
|
function exposureDecision(definition, { role = "user", user = {}, context = null } = {}) {
|
|
const tool = publicDefinition(definition);
|
|
if (!roleAllows(role, definition.required_role)) {
|
|
return hidden(tool, "permission_blocked", "Permission denied: the selected role is below the tool's required role.");
|
|
}
|
|
if (!definition.read_only && context?.permission_context?.webui_actions_allowed === false) {
|
|
return hidden(tool, "scope_blocked", "This action tool requires a WebUI action-capable context.");
|
|
}
|
|
const origin = contextOrigin(context);
|
|
if (Array.isArray(definition.allowed_origins) && !definition.allowed_origins.includes(origin)) {
|
|
return hidden(tool, "origin_blocked", `The tool does not allow the ${origin} origin.`);
|
|
}
|
|
if (typeof definition.origin_check === "function") {
|
|
try {
|
|
const result = definition.origin_check({ user, context, origin, phase: "prompt" });
|
|
if (result && typeof result.then === "function") {
|
|
return hidden(tool, "origin_blocked", "Tool origin checks must be synchronous.");
|
|
}
|
|
if (result !== true) return hidden(tool, "origin_blocked", `The tool is disabled for the ${origin} origin.`);
|
|
} catch {
|
|
return hidden(tool, "origin_blocked", "The tool origin check failed.");
|
|
}
|
|
}
|
|
if (typeof definition.scope_check === "function") {
|
|
try {
|
|
const result = definition.scope_check({ user, context, origin, phase: "prompt" });
|
|
if (result && typeof result.then === "function") {
|
|
return hidden(tool, "scope_blocked", "Tool scope checks must be synchronous.");
|
|
}
|
|
if (result !== true) return hidden(tool, "scope_blocked", "The tool is outside this request context.");
|
|
} catch {
|
|
return hidden(tool, "scope_blocked", "The tool scope check failed.");
|
|
}
|
|
}
|
|
const promptCheck = definition.prompt_permission_check || definition.permission_check;
|
|
try {
|
|
const allowed = promptCheck({
|
|
user,
|
|
arguments: {},
|
|
required_permission: definition.required_permission,
|
|
context,
|
|
phase: "prompt"
|
|
});
|
|
if (allowed && typeof allowed.then === "function") {
|
|
return hidden(tool, "permission_blocked", "Tool prompt permission checks must be synchronous.");
|
|
}
|
|
if (allowed !== true) return hidden(tool, "permission_blocked", "The backend permission check denied prompt exposure.");
|
|
} catch {
|
|
return hidden(tool, "permission_blocked", "The backend permission check denied prompt exposure.");
|
|
}
|
|
return { tool, exposed: true, reason: "exposed", message: "Registered and permitted for this context." };
|
|
}
|
|
|
|
function parseToolCall(text) {
|
|
const parsed = parseToolCallResult(text);
|
|
return parsed.status === "valid" ? parsed.call : null;
|
|
}
|
|
|
|
function parseToolCallResult(text) {
|
|
const value = String(text || "").trim();
|
|
if (!value) return { status: "none", call: null, error: null };
|
|
let parsed;
|
|
try {
|
|
parsed = JSON.parse(value);
|
|
} catch {
|
|
return /"type"\s*:\s*"tool_call"/i.test(value)
|
|
? { status: "malformed", call: null, error: "Tool-call JSON was malformed or surrounded by extra text." }
|
|
: { status: "none", call: null, error: null };
|
|
}
|
|
if (parsed?.type !== "tool_call") return { status: "none", call: null, error: null };
|
|
const keys = Object.keys(parsed);
|
|
if (keys.length !== 3 || !["type", "tool", "arguments"].every((key) => keys.includes(key))) {
|
|
return { status: "malformed", call: null, error: "Tool-call JSON must contain only type, tool, and arguments." };
|
|
}
|
|
if (typeof parsed.tool !== "string" || !parsed.tool || !parsed.arguments ||
|
|
typeof parsed.arguments !== "object" || Array.isArray(parsed.arguments)) {
|
|
return { status: "malformed", call: null, error: "Tool-call JSON is missing a valid tool or arguments object." };
|
|
}
|
|
return {
|
|
status: "valid",
|
|
call: { type: "tool_call", tool: parsed.tool, arguments: parsed.arguments },
|
|
error: null
|
|
};
|
|
}
|
|
|
|
function publicDefinition(definition) {
|
|
const {
|
|
workflow_handler,
|
|
permission_check,
|
|
prompt_permission_check,
|
|
origin_check,
|
|
scope_check,
|
|
...tool
|
|
} = definition;
|
|
return tool;
|
|
}
|
|
|
|
function validateDefinition(definition) {
|
|
if (!definition?.tool_id || !definition.display_name || !definition.description ||
|
|
!definition.owning_plugin || !definition.required_permission || !definition.audit_category ||
|
|
typeof definition.workflow_handler !== "function" ||
|
|
typeof definition.permission_check !== "function" ||
|
|
!definition.schema || typeof definition.schema !== "object" || Array.isArray(definition.schema)) {
|
|
throw new Error("Invalid AI tool definition.");
|
|
}
|
|
}
|
|
|
|
function normalizeArgument(key, value, descriptor) {
|
|
if (descriptor.type === "integer") {
|
|
if (!Number.isInteger(Number(value))) throw toolError("schema_invalid", `${key} must be an integer.`);
|
|
value = Number(value);
|
|
} else if (descriptor.type === "number") {
|
|
if (!Number.isFinite(Number(value))) throw toolError("schema_invalid", `${key} must be a number.`);
|
|
value = Number(value);
|
|
} else if (descriptor.type === "string") {
|
|
if (typeof value !== "string") throw toolError("schema_invalid", `${key} must be a string.`);
|
|
} else if (descriptor.type === "boolean") {
|
|
if (typeof value !== "boolean") throw toolError("schema_invalid", `${key} must be a boolean.`);
|
|
} else if (descriptor.type === "array") {
|
|
if (!Array.isArray(value)) throw toolError("schema_invalid", `${key} must be an array.`);
|
|
} else if (descriptor.type === "object") {
|
|
if (!value || typeof value !== "object" || Array.isArray(value)) {
|
|
throw toolError("schema_invalid", `${key} must be an object.`);
|
|
}
|
|
}
|
|
if (Array.isArray(descriptor.enum) && !descriptor.enum.includes(value)) {
|
|
throw toolError("schema_invalid", `${key} is invalid.`);
|
|
}
|
|
return value;
|
|
}
|
|
|
|
function summarizeArguments(args) {
|
|
return Object.fromEntries(Object.entries(args || {}).map(([key, value]) => {
|
|
if (/token|secret|password|key|authorization|cookie/i.test(key)) return [key, "[REDACTED]"];
|
|
if (typeof value === "string") return [key, { type: "string", length: value.length }];
|
|
if (Array.isArray(value)) return [key, { type: "array", length: value.length }];
|
|
return [key, { type: typeof value }];
|
|
}));
|
|
}
|
|
|
|
function contextOrigin(context) {
|
|
const value = String(context?.origin || context?.platform || "other").toLowerCase();
|
|
return ["webui", "discord", "twitch", "youtube", "kick"].includes(value) ? value : "other";
|
|
}
|
|
|
|
function auditActor(user, role = null) {
|
|
return {
|
|
user_id: user?.id || null,
|
|
username: user?.username || null,
|
|
role: role || (user?.isAdmin ? "admin" : user?.isMod ? "mod" : "user")
|
|
};
|
|
}
|
|
|
|
function hidden(tool, reason, message) {
|
|
return { tool, exposed: false, reason, message };
|
|
}
|
|
|
|
function toolError(code, message) {
|
|
return Object.assign(new Error(message), { code });
|
|
}
|
|
|
|
module.exports = {
|
|
ToolRegistry,
|
|
contextOrigin,
|
|
exposureDecision,
|
|
parseToolCall,
|
|
parseToolCallResult,
|
|
publicDefinition,
|
|
summarizeArguments
|
|
};
|