Executing Tools Safely: Validation and Timeouts

type JsonPrimitive = string | number | boolean | null;
export type JsonValue = JsonPrimitive | JsonValue[] | { [k: string]: JsonValue };
type ValidationIssue = {
  path: string;
  message: string;
};
export type ToolCallRecord = {
  toolCallId: string;
  name: string;
  args: JsonValue;
  startedAt: string;
  durationMs: number;
};
export type ToolSuccess = {
  kind: "success";
  record: ToolCallRecord;
  output: JsonValue;
};
export type ToolValidationError = {
  kind: "validation_error";
  record: ToolCallRecord;
  issues: ValidationIssue[];
};
export type ToolExecutionError = {
  kind: "execution_error";
  record: ToolCallRecord;
  error: {
    name: string;
    message: string;
    stack?: string;
    code?: string;
  };
};
export type ToolResult = ToolSuccess | ToolValidationError | ToolExecutionError;
export type ToolContext = {
  toolCallId: string;
  startedAtMs: number;
  timeoutMs: number;
};
export type ToolSpec = {
  name: string;
  description: string;
  validate: (args: JsonValue) => ValidationIssue[];
  run: (args: JsonValue, ctx: ToolContext) => JsonValue | Promise<JsonValue>;
};
const registry = new Map<string, ToolSpec>();
function registerTool(spec: ToolSpec): void {
  // Enforce unique names in the registry.
  if (registry.has(spec.name)) {
    throw new Error(`Duplicate tool name: ${spec.name}`);
  }
  registry.set(spec.name, spec);
}
function nowMs(): number {
  // Use epoch ms for duration calculations.
  return Date.now();
}
function createToolCallId(): string {
  // Avoid external deps; ensure serializable id.
  return `tc_${nowMs()}_${Math.random().toString(16).slice(2)}`;
}
function asErrorShape(err: unknown): ToolExecutionError["error"] {
  // Normalize unknown throws into JSON.
  if (err instanceof Error) {
    return {
      name: err.name || "Error",
      message: err.message || String(err),
      stack: err.stack,
    };
  }
  return {
    name: "ThrownValue",
    message: typeof err === "string" ? err : JSON.stringify(err),
  };
}
function withTimeout<T>(promise: Promise<T>, timeoutMs: number): Promise<T> {
  // Use Promise.race to enforce timeouts.
  return new Promise<T>((resolve, reject) => {
    const t = setTimeout(() => {
      const e = new Error(`Tool timed out after ${timeoutMs}ms`);
      (e as any).code = "TIMEOUT";
      reject(e);
    }, timeoutMs);
    promise
      .then((v) => {
        clearTimeout(t);
        resolve(v);
      })
      .catch((e) => {
        clearTimeout(t);
        reject(e);
      });
  });
}
export async function runTool(input: {
  name: string;
  args: JsonValue;
  timeoutMs?: number;
}): Promise<ToolResult> {
  // Treat args as untrusted even if parsed.
  const toolCallId = createToolCallId();
  const startedAtMs = nowMs();
  const startedAt = new Date(startedAtMs).toISOString();
  const timeoutMs = input.timeoutMs ?? 2_000;
  const baseRecord: ToolCallRecord = {
    toolCallId,
    name: input.name,
    args: input.args,
    startedAt,
    durationMs: 0,
  };
  const spec = registry.get(input.name);
  if (!spec) {
    // Unknown tool becomes a validation error.
    const endedAtMs = nowMs();
    return {
      kind: "validation_error",
      record: { ...baseRecord, durationMs: endedAtMs - startedAtMs },
      issues: [{ path: "name", message: `Unknown tool: ${input.name}` }],
    };
  }
  const issues = spec.validate(input.args);
  if (issues.length > 0) {
    // Validation errors are returned, not thrown.
    const endedAtMs = nowMs();
    return {
      kind: "validation_error",
      record: { ...baseRecord, durationMs: endedAtMs - startedAtMs },
      issues,
    };
  }
  const ctx: ToolContext = { toolCallId, startedAtMs, timeoutMs };
  try {
    // Run tool with enforced timeout.
    const output = await withTimeout(Promise.resolve(spec.run(input.args, ctx)), timeoutMs);
    const endedAtMs = nowMs();
    return {
      kind: "success",
      record: { ...baseRecord, durationMs: endedAtMs - startedAtMs },
      output,
    };
  } catch (err) {
    // Convert thrown values into JSON-safe shape.
    const endedAtMs = nowMs();
    const shaped = asErrorShape(err);
    const code = (err as any)?.code;
    return {
      kind: "execution_error",
      record: { ...baseRecord, durationMs: endedAtMs - startedAtMs },
      error: {
        ...shaped,
        code: typeof code === "string" ? code : undefined,
      },
    };
  }
}
function isPlainObject(v: unknown): v is Record<string, unknown> {
  // Ensure args have expected container type.
  return typeof v === "object" && v !== null && !Array.isArray(v);
}
function validateAddNumbersArgs(args: JsonValue): ValidationIssue[] {
  // Validate shape and types for untrusted input.
  const issues: ValidationIssue[] = [];
  if (!isPlainObject(args)) {
    issues.push({ path: "", message: "Expected an object" });
    return issues;
  }
  const a = (args as any).a;
  const b = (args as any).b;
  if (typeof a !== "number" || Number.isNaN(a)) {
    issues.push({ path: "a", message: "Expected a finite number" });
  }
  if (typeof b !== "number" || Number.isNaN(b)) {
    issues.push({ path: "b", message: "Expected a finite number" });
  }
  return issues;
}
function validateDelayEchoArgs(args: JsonValue): ValidationIssue[] {
  // Validate delay bounds and message type.
  const issues: ValidationIssue[] = [];
  if (!isPlainObject(args)) {
    issues.push({ path: "", message: "Expected an object" });
    return issues;
  }
  const ms = (args as any).ms;
  const message = (args as any).message;
  if (typeof ms !== "number" || !Number.isFinite(ms) || ms < 0 || ms > 10_000) {
    issues.push({ path: "ms", message: "Expected number 0..10000" });
  }
  if (typeof message !== "string") {
    issues.push({ path: "message", message: "Expected a string" });
  }
  return issues;
}
function sleep(ms: number): Promise<void> {
  // Simple timer wrapper for delay tool.
  return new Promise((resolve) => setTimeout(resolve, ms));
}
registerTool({
  name: "addNumbers",
  description: "Adds two numbers: { a, b }",
  validate: validateAddNumbersArgs,
  run: (args) => {
    // Cast only after validation succeeds.
    const { a, b } = args as any as { a: number; b: number };
    return { sum: a + b };
  },
});
registerTool({
  name: "delayEcho",
  description: "Waits ms then echoes message: { ms, message }",
  validate: validateDelayEchoArgs,
  run: async (args, ctx) => {
    // Avoid trusting delay values beyond validator.
    const { ms, message } = args as any as { ms: number; message: string };
    await sleep(ms);
    return {
      echoed: message,
      waitedMs: ms,
      timeoutMs: ctx.timeoutMs,
    };
  },
});
registerTool({
  name: "explode",
  description: "Throws to demonstrate execution_error",
  validate: (args) => {
    // Require a reason string for clarity.
    if (!isPlainObject(args) || typeof (args as any).reason !== "string") {
      return [{ path: "reason", message: "Expected a string" }];
    }
    return [];
  },
  run: (args) => {
    // Intentionally throw to show error handling.
    const { reason } = args as any as { reason: string };
    throw new Error(`Boom: ${reason}`);
  },
});
(async () => {
  const calls: Array<{ name: string; args: JsonValue; timeoutMs?: number; label: string }> = [
    { label: "success", name: "addNumbers", args: { a: 2, b: 40 } },
    { label: "validation_error", name: "addNumbers", args: { a: "2", b: 1 } as any },
    { label: "unknown_tool", name: "missingTool", args: { any: "thing" } },
    { label: "execution_error", name: "explode", args: { reason: "demo" } },
    { label: "timeout", name: "delayEcho", args: { ms: 1500, message: "hi" }, timeoutMs: 100 },
  ];
  for (const c of calls) {
    // Print each result as standalone JSON.
    console.log("-".repeat(66));
    console.log(`Case: ${c.label}`);
    const res = await runTool({ name: c.name, args: c.args, timeoutMs: c.timeoutMs });
    console.log(JSON.stringify(res, null, 2));
  }
  console.log("-".repeat(66));
})();