9 files changed, 32 insertions, 24 deletions
diff --git a/scripts/build-declarations.js b/scripts/build-declarations.js
index e1a2892..35a75ee 100644
--- a/scripts/build-declarations.js
+++ b/scripts/build-declarations.js
@@ -2,19 +2,23 @@
 
 const fs = require("fs");
 
-function main() {
+const TOOL_ENTRY_FUNC = "run";
+
+function main(isTool = true) {
   const scriptfile = process.argv[2];
   const contents = fs.readFileSync(process.argv[2], "utf8");
-  const functions = extractFunctions(contents);
+  const functions = extractFunctions(contents, isTool);
   let declarations = functions.map(({ funcName, jsdoc }) => {
     const { description, params } = parseJsDoc(jsdoc, funcName);
     const declaration = buildDeclaration(funcName, description, params);
     return declaration;
   });
-  const name = getBasename(scriptfile);
-  if (declarations.length > 0) {
-    declarations = declarations.slice(0, 1);
-    declarations[0].name = name;
+  if (isTool) {
+    const name = getBasename(scriptfile);
+    if (declarations.length > 0) {
+      declarations = declarations.slice(0, 1);
+      declarations[0].name = name;
+    }
   }
   console.log(JSON.stringify(declarations, null, 2));
 }
@@ -23,7 +27,7 @@ function main() {
  * @param {string} contents
  * @param {bool} isTool
  */
-function extractFunctions(contents, isTool = true) {
+function extractFunctions(contents, isTool) {
   const output = [];
   const lines = contents.split("\n");
   let isInComment = false;
@@ -45,9 +49,9 @@ function extractFunctions(contents, isTool = true) {
         continue;
       }
       if (isTool) {
-        if (/function main/.test(line)) {
+        if (new RegExp(`function ${TOOL_ENTRY_FUNC}`).test(line)) {
           output.push({
-            funcName: "main",
+            funcName: TOOL_ENTRY_FUNC,
             jsdoc,
           });
         }
diff --git a/scripts/build-declarations.py b/scripts/build-declarations.py
index 17a27b6..2a28e43 100644
--- a/scripts/build-declarations.py
+++ b/scripts/build-declarations.py
@@ -7,13 +7,14 @@ import re
 import sys
 from collections import OrderedDict
 
+TOOL_ENTRY_FUNC = "run"
 
-def main():
+def main(is_tool = True):
     scriptfile = sys.argv[1]
     with open(scriptfile, "r", encoding="utf-8") as f:
         contents = f.read()
 
-    functions = extract_functions(contents)
+    functions = extract_functions(contents, is_tool)
     declarations = []
     for function in functions:
         func_name, docstring, func_args = function
@@ -22,15 +23,16 @@ def main():
             build_declaration(func_name, description, params, func_args)
         )
 
-    name = os.path.splitext(os.path.basename(scriptfile))[0]
-    if declarations:
-        declarations = declarations[0:1]
-        declarations[0]["name"] = name
+    if is_tool:
+        name = os.path.splitext(os.path.basename(scriptfile))[0]
+        if declarations:
+            declarations = declarations[0:1]
+            declarations[0]["name"] = name
 
     print(json.dumps(declarations, indent=2))
 
 
-def extract_functions(contents: str):
+def extract_functions(contents: str, is_tool: bool):
     tree = ast.parse(contents)
     output = []
     for node in ast.walk(tree):
@@ -39,6 +41,8 @@ def extract_functions(contents: str):
         func_name = node.name
         if func_name.startswith("_"):
             continue
+        if is_tool and func_name != TOOL_ENTRY_FUNC:
+            continue
         docstring = ast.get_docstring(node) or ""
         func_args = OrderedDict()
         for arg in node.args.args:
diff --git a/scripts/create-tool.sh b/scripts/create-tool.sh
index b6d56dc..ff0ce5e 100755
--- a/scripts/create-tool.sh
+++ b/scripts/create-tool.sh
@@ -75,7 +75,7 @@ create_js() {
  * @typedef {Object} Args${properties}
  * @param {Args} args
  */
-exports.main = function main(args) {
+exports.run = function run(args) {
   console.log(args);
 }
 EOF
diff --git a/scripts/run-tool.js b/scripts/run-tool.js
index f99b2e0..d0cafa4 100755
--- a/scripts/run-tool.js
+++ b/scripts/run-tool.js
@@ -71,5 +71,5 @@ try {
   process.exit(1);
 }
 
-const { main } = loadFunc(funcName);
-main(args);
+const { run } = loadFunc(funcName);
+run(args);
diff --git a/scripts/run-tool.py b/scripts/run-tool.py
index 97d1d25..c6bcfc8 100755
--- a/scripts/run-tool.py
+++ b/scripts/run-tool.py
@@ -72,4 +72,4 @@ except (json.JSONDecodeError, TypeError):
     sys.exit(1)
 
 module = load_func(func_name)
-module.main(**args)
+module.run(**args)
diff --git a/tools/demo_tool.js b/tools/demo_tool.js
index f8fd283..d47fa59 100644
--- a/tools/demo_tool.js
+++ b/tools/demo_tool.js
@@ -11,6 +11,6 @@
  * @property {string[]} [array_optional] - Define a optional string array property
  * @param {Args} args
  */
-exports.main = function main(args) {
+exports.run = function run(args) {
   console.log(args);
 }
diff --git a/tools/demo_tool.py b/tools/demo_tool.py
index bd353ad..5eb4c09 100644
--- a/tools/demo_tool.py
+++ b/tools/demo_tool.py
@@ -1,7 +1,7 @@
 from typing import List, Literal, Optional
 
 
-def main(
+def run(
     boolean: bool,
     string: str,
     string_enum: Literal["foo", "bar"],
diff --git a/tools/may_execute_js_code.js b/tools/may_execute_js_code.js
index 2dec177..4706e07 100644
--- a/tools/may_execute_js_code.js
+++ b/tools/may_execute_js_code.js
@@ -4,6 +4,6 @@
  * @property {string} code - Javascript code to execute, such as `console.log("hello world")`
  * @param {Args} args
  */
-exports.main = function main({ code }) {
+exports.run = function run({ code }) {
   eval(code);
 }
diff --git a/tools/may_execute_py_code.py b/tools/may_execute_py_code.py
index 8e55a8d..5f6af2f 100644
--- a/tools/may_execute_py_code.py
+++ b/tools/may_execute_py_code.py
@@ -1,4 +1,4 @@
-def main(code: str):
+def run(code: str):
     """Runs the python code.
     Args:
         code: Python code to execute, such as `print("hello world")`