From 1bcb6aa36681d4f2fe6c6cdf4d51999708df6528 Mon Sep 17 00:00:00 2001 From: Jiongan Mu Date: Wed, 15 Jul 2026 00:25:05 -0700 Subject: [PATCH] add node-embedded custom tools --- README.md | 26 ++++ backend/app/api/call_chats.py | 11 ++ backend/app/orchestrator/prompt.py | 19 ++- backend/app/runner/chat.py | 8 +- backend/app/runner/chat_child.py | 12 +- backend/app/runner/child.py | 5 +- backend/app/runner/ctx.py | 149 +++++++++++++++++++- backend/app/runner/llm.py | 22 ++- backend/app/runner/node_tools.py | 67 +++++++++ backend/tests/test_call_chats.py | 10 +- backend/tests/test_node_tools.py | 209 +++++++++++++++++++++++++++++ backend/tests/test_orchestrator.py | 9 ++ backend/tests/test_runner.py | 39 ++++++ 13 files changed, 566 insertions(+), 20 deletions(-) create mode 100644 backend/app/runner/node_tools.py create mode 100644 backend/tests/test_node_tools.py diff --git a/README.md b/README.md index 85f99c2..94ba95e 100644 --- a/README.md +++ b/README.md @@ -94,6 +94,32 @@ def run(inputs, ctx): The six built-in tools available to every node are `shell`, `read_file`, `write_file`, `edit_file`, `web_search`, and `web_fetch`. `read_file` also returns images (PNG/JPEG/GIF/WebP) as attachments for vision-capable models; `web_search` and `web_fetch` are backed by parallel.ai. Tools can be invoked either *agentically* — named in `ctx.agent(tools=[...])` so the node's own model decides when to call them — or *directly* via `ctx.tools.(...)`, which runs them deterministically with no model in the loop (see [Design decisions](#design-decisions)). +Nodes can also declare custom tools inside the same Python block—there is no separate node field or registry configuration. A lowercase `NodeTool` subclass registers automatically while the node source is loaded, and then uses the same quoted-name surfaces as built-ins: + +```python +from app.runner.node_tools import NodeTool + +class lookup_order(NodeTool): + description = "Look up an order by id." + parameters = { + "type": "object", + "properties": {"order_id": {"type": "string"}}, + "required": ["order_id"], + } + + def execute(self, ctx, order_id): + return {"order_id": order_id, "status": "shipped"} + +def run(inputs, ctx): + result = ctx.agent( + prompt=f"Help with order {inputs['order_id']}", + tools=["lookup_order"], + ) + return {"answer": result["content"]} +``` + +Custom tool names must use the same lowercase `snake_case` convention as native tools and cannot shadow built-in or MCP names. They are isolated to their node, included automatically in frozen run source, traced like native calls, and restored from that frozen source when an agent conversation is continued. + ## Design decisions A few choices that shape how the system behaves in use: diff --git a/backend/app/api/call_chats.py b/backend/app/api/call_chats.py index d4df862..425f97c 100644 --- a/backend/app/api/call_chats.py +++ b/backend/app/api/call_chats.py @@ -86,6 +86,16 @@ def _node_name(db: Session, node_run: models.NodeRun) -> str: return node_run.node_id +def _snapshot_node_code(db: Session, node_run: models.NodeRun) -> str: + """Frozen source for the node that produced ``node_run``.""" + + run = db.get(models.Run, node_run.run_id) + for node in ((run.workflow_snapshot if run else None) or {}).get("nodes", []) or []: + if node.get("id") == node_run.node_id: + return node.get("code") or "" + return "" + + def _build_continuation(db: Session, nr: models.NodeRun, call_id: str) -> models.CallChat: """Build — but do NOT persist — the continuation for one agent call, seeded from its recorded transcript. Raises 404 if the node run, the call, @@ -255,6 +265,7 @@ def send_call_chat_turn( tools=list(chat.tools or []), model=model, child_env=child_env, + node_code=_snapshot_node_code(db, nr), ) return schemas.CallChatTurnOut(turn_id=turn_id) diff --git a/backend/app/orchestrator/prompt.py b/backend/app/orchestrator/prompt.py index f6858de..b7c38cd 100644 --- a/backend/app/orchestrator/prompt.py +++ b/backend/app/orchestrator/prompt.py @@ -159,6 +159,23 @@ def _format_node_tool_names() -> str: Every node defines a `run(inputs, ctx)` function. Top-level `import`s and small helper functions alongside `run` are fine — the whole code blob is `exec`'d into a fresh namespace per run, so reach for `json`, `re`, `pathlib`, etc. when they're cleaner than routing through an LLM. +A node may also declare custom tools directly in the same code blob. Subclass `NodeTool`; the lowercase snake_case class name is the tool name and the runtime registers it automatically when the node source executes. Do not add a separate registry, decorator, or list. Custom tools are node-local, may call `ctx.tools.*`, and are available through the same name-based surfaces as native tools: agentically (`tools=["lookup_order"]`) and directly (`ctx.tools.lookup_order(...)`). Keep top-level code declarative (imports, helpers, tool classes, and `run`) because the frozen source is also loaded when continuing one of the node's agent calls. + +```python +from app.runner.node_tools import NodeTool + +class lookup_order(NodeTool): + description = "Look up an order by id." + parameters = { + "type": "object", + "properties": {"order_id": {"type": "string"}}, + "required": ["order_id"], + } + + def execute(self, ctx, order_id): + return {"order_id": order_id, "status": "shipped"} +``` + If a node imports a third-party package, make sure it's installed *before* the workflow runs. ```python @@ -171,7 +188,7 @@ def run(inputs, ctx): `ctx` provides: -- `ctx.agent(prompt, tools=[...])` — runs an LLM inside the node. Pass tool names ([[NODE_TOOL_NAMES]]) in the `tools` list; the LLM running inside the node decides when to invoke them. Returns a dict with keys `content` (str), `tool_calls_made` (list), `usage`, `cost`. Omit the `model` arg (see *# design conventions*). Optional `label` when a node makes several calls — keep it short and meaningful. +- `ctx.agent(prompt, tools=[...])` — runs an LLM inside the node. Pass registered tool names ([[NODE_TOOL_NAMES]] plus any lowercase custom `NodeTool` class names) as strings; the LLM running inside the node decides when to invoke them. Returns a dict with keys `content` (str), `tool_calls_made` (list), `usage`, `cost`. Omit the `model` arg (see *# design conventions*). Optional `label` when a node makes several calls — keep it short and meaningful. - `ctx.tools.shell(...)` / `ctx.tools.read_file(...)` / `ctx.tools.web_fetch(...)` / … — direct (non-LLM) tool calls, same names, returning the same dicts the LLM-mediated form would produce. The agentic form above is the default; reserve direct calls for when there's nothing for a model to decide (see *# direct calls vs wrapping the tool in an agent*). - `ctx.log("...")` — appends a visible line to the run log. - `ctx.workdir` — `pathlib.Path` to a per-run scratch directory. diff --git a/backend/app/runner/chat.py b/backend/app/runner/chat.py index 46d1c90..eb60556 100644 --- a/backend/app/runner/chat.py +++ b/backend/app/runner/chat.py @@ -176,6 +176,7 @@ def start_chat_turn( tools: list[str], model: str, child_env: dict[str, str], + node_code: str = "", ) -> None: """Begin a chat turn in the background. Returns immediately. @@ -189,7 +190,7 @@ def start_chat_turn( ev_mod.get_or_create(turn_id) threading.Thread( target=_run_turn, - args=(turn_id, chat_id, messages, tools, model, child_env), + args=(turn_id, chat_id, messages, tools, model, child_env, node_code), daemon=True, ).start() @@ -215,6 +216,7 @@ def _run_turn( tools: list[str], model: str, child_env: dict[str, str], + node_code: str = "", ) -> None: workdir = tempfile.mkdtemp(prefix="wfchat-") try: @@ -222,6 +224,10 @@ def _run_turn( "messages": messages, "tools": tools, "model": model, + # Frozen source from the run that produced this conversation. + # The chat child re-executes it to recover embedded NodeTool + # subclasses without storing a second copy on CallChat. + "node_code": node_code, "workdir": workdir, "env": child_env, } diff --git a/backend/app/runner/chat_child.py b/backend/app/runner/chat_child.py index 00f021d..c72479c 100644 --- a/backend/app/runner/chat_child.py +++ b/backend/app/runner/chat_child.py @@ -7,7 +7,8 @@ Reads a JSON payload from stdin:: - {"messages": [...], "tools": [...], "model": "...", "workdir": "...", "env": {...}} + {"messages": [...], "tools": [...], "model": "...", "node_code": "...", + "workdir": "...", "env": {...}} Emits the same per-call event contract a run does — ``llm_call_started``, ``llm_round_started``, ``llm_call_chunk``, ``tool_call_started/finished``, @@ -36,6 +37,7 @@ def main() -> None: payload = _read_payload() messages = payload.get("messages") or [] tools = payload.get("tools") or [] + node_code = payload.get("node_code") or "" model = payload.get("model") or payload.get("default_model") or "" workdir = Path(payload.get("workdir") or ".") workdir.mkdir(parents=True, exist_ok=True) @@ -55,6 +57,14 @@ def main() -> None: cancelled = False try: + # A continuation exposes the exact custom tools from the run it + # continues, not whatever the live node contains now. The API sends + # frozen node source; loading it registers NodeTool subclasses but + # never invokes the node's ``run`` function. + if node_code: + from app.runner.node_tools import execute_node_source + + execute_node_source(node_code, ctx.register_node_tool) result = ctx.agent(model=model, prompt=messages, tools=tools) except KeyboardInterrupt: cancelled = True diff --git a/backend/app/runner/child.py b/backend/app/runner/child.py index 3043b9b..a3851a4 100644 --- a/backend/app/runner/child.py +++ b/backend/app/runner/child.py @@ -141,8 +141,9 @@ def _emit_node_skipped(node_id: str, inputs: dict, output_ports: list[dict]) -> def _execute_node(node: dict, ctx, inputs: dict) -> dict: """Exec the node's user code and return its (port-normalised) outputs. Raises if the code is malformed or returns the wrong shape.""" - ns: dict = {} - exec(node.get("code") or "", ns, ns) + from app.runner.node_tools import execute_node_source + + ns = execute_node_source(node.get("code") or "", ctx.register_node_tool) run_fn = ns.get("run") if not callable(run_fn): raise RuntimeError("node code must define `run(inputs, ctx)` function") diff --git a/backend/app/runner/ctx.py b/backend/app/runner/ctx.py index 0ce49e1..5e0ba55 100644 --- a/backend/app/runner/ctx.py +++ b/backend/app/runner/ctx.py @@ -11,6 +11,8 @@ from __future__ import annotations import inspect import itertools +import json +import re import os import threading from pathlib import Path @@ -19,10 +21,12 @@ from app.runner.tools import ( MCP_NAMESPACES, REGISTRY, + TOOL_SCHEMAS, mcp_unavailable_error, strip_attachment_data, ) from app.runner import llm as llm_mod +from app.runner.node_tools import NodeTool EmitFn = Callable[[dict], None] @@ -67,7 +71,14 @@ class _ToolsProxy: (``ctx.tools.notion_create_pages(...)``) and dotted by server (``ctx.tools.notion.create_pages(...)``).""" - def __init__(self, recorder: list[dict], on_event: EmitFn, lock: threading.Lock): + def __init__( + self, + registry: dict[str, Callable], + recorder: list[dict], + on_event: EmitFn, + lock: threading.Lock, + ): + self._registry = registry self._recorder = recorder self._on_event = on_event self._lock = lock @@ -77,7 +88,7 @@ def __init__(self, recorder: list[dict], on_event: EmitFn, lock: threading.Lock) self._call_counter = itertools.count(1) def __getattr__(self, name: str): - if name in REGISTRY: + if name in self._registry: return self._bind(name) if name in MCP_NAMESPACES: return _ServerProxy(self, name, MCP_NAMESPACES[name]) @@ -89,7 +100,7 @@ def __getattr__(self, name: str): raise AttributeError(f"no tool '{name}' in registry") def _bind(self, name: str): - fn = REGISTRY.get(name) + fn = self._registry.get(name) if fn is None: raise AttributeError(f"no tool '{name}' in registry") @@ -104,6 +115,9 @@ def wrapped(*args, **kwargs): bound = sig.bind(*args, **kwargs) bound.apply_defaults() call_args = dict(bound.arguments) + for param in sig.parameters.values(): + if param.kind == inspect.Parameter.VAR_KEYWORD: + call_args.update(call_args.pop(param.name, {})) tc_id = f"direct-{next(self._call_counter)}" self._on_event( @@ -171,7 +185,122 @@ def __init__( self.tool_calls: list[dict] = [] self._lock = threading.Lock() self._call_counter = itertools.count(1) - self.tools = _ToolsProxy(self.tool_calls, self._on_event, self._lock) + # MCP discovery has already populated the process registry before a + # node Ctx is created. Copy it here so custom tools stay node-local: + # parallel nodes may use the same custom name without colliding. + self.tool_registry: dict[str, Callable] = dict(REGISTRY) + self.tool_schemas: dict[str, dict] = dict(TOOL_SCHEMAS) + self.tools = _ToolsProxy( + self.tool_registry, self.tool_calls, self._on_event, self._lock + ) + + def register_node_tool(self, tool_class: type[NodeTool]) -> None: + """Validate and register one embedded ``NodeTool`` subclass.""" + + if not inspect.isclass(tool_class) or not issubclass(tool_class, NodeTool): + raise TypeError("custom tools must subclass NodeTool") + + name = tool_class.__name__ + if not re.fullmatch(r"[a-z][a-z0-9_]*", name): + raise ValueError( + f"custom tool class '{name}' must use lowercase snake_case" + ) + if name in self.tool_registry: + raise ValueError(f"custom tool '{name}' conflicts with an existing tool") + + description = getattr(tool_class, "description", "") + if not isinstance(description, str) or not description.strip(): + raise ValueError(f"custom tool '{name}' must define a description") + + parameters = getattr(tool_class, "parameters", None) + if not isinstance(parameters, dict) or parameters.get("type") != "object": + raise ValueError( + f"custom tool '{name}' parameters must be an object JSON schema" + ) + properties = parameters.get("properties", {}) + required = parameters.get("required", []) + if not isinstance(properties, dict): + raise ValueError( + f"custom tool '{name}' parameters.properties must be an object" + ) + if not isinstance(required, list) or any( + not isinstance(item, str) or item not in properties for item in required + ): + raise ValueError( + f"custom tool '{name}' parameters.required must name declared properties" + ) + if "execute" not in tool_class.__dict__ or not callable(tool_class.execute): + raise ValueError(f"custom tool '{name}' must define execute(self, ctx, ...)") + try: + json.dumps(parameters) + except (TypeError, ValueError) as e: + raise ValueError( + f"custom tool '{name}' parameters must be JSON-serializable: {e}" + ) from e + + execute_params = list(inspect.signature(tool_class.execute).parameters.values()) + if ( + len(execute_params) < 2 + or execute_params[0].name != "self" + or execute_params[1].name != "ctx" + ): + raise ValueError( + f"custom tool '{name}' execute signature must start with (self, ctx)" + ) + tool_params = execute_params[2:] + if any( + p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.VAR_POSITIONAL) + for p in tool_params + ): + raise ValueError( + f"custom tool '{name}' execute arguments must be named parameters" + ) + has_var_kwargs = any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in tool_params + ) + if not has_var_kwargs: + declared = { + p.name + for p in tool_params + if p.kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + } + signature_required = { + p.name + for p in tool_params + if p.name in declared and p.default is inspect.Parameter.empty + } + if declared != set(properties): + raise ValueError( + f"custom tool '{name}' schema properties must match its execute arguments" + ) + if signature_required != set(required): + raise ValueError( + f"custom tool '{name}' schema required fields must match " + "execute arguments without defaults" + ) + + def handler(*args, **arguments): + return tool_class().execute(self, *args, **arguments) + + handler.__name__ = name + # Make the wrapper look exactly like execute minus its injected + # (self, ctx), so direct calls get native-style argument validation and + # flat trace payloads while LLM dispatch can continue calling by kwargs. + handler.__signature__ = inspect.Signature(tool_params) # type: ignore[attr-defined] + + self.tool_registry[name] = handler + self.tool_schemas[name] = { + "type": "function", + "function": { + "name": name, + "description": description.strip(), + "parameters": parameters, + }, + } def _next_call_id(self) -> str: return f"call-{next(self._call_counter)}" @@ -194,13 +323,17 @@ def agent( if not m: raise RuntimeError("agent: no model specified and no default configured") + tool_names = list(tools or []) + if any(not isinstance(name, str) for name in tool_names): + raise TypeError("agent tools must be registered tool-name strings") + call_id = self._next_call_id() call_label = (label or "").strip() or None started: dict = { "type": "llm_call_started", "call_id": call_id, "model": m, - "tools": tools or [], + "tools": tool_names, } if call_label: started["label"] = call_label @@ -209,7 +342,9 @@ def agent( result = llm_mod.call_llm( m, prompt, - tools=tools, + tools=tool_names, + tool_registry=self.tool_registry, + tool_schemas_by_name=self.tool_schemas, on_event=self._on_event, call_id=call_id, **opts, @@ -236,7 +371,7 @@ def agent( "call_id": call_id, "model": m, "prompt": prompt if isinstance(prompt, str) else "", - "tools": tools or [], + "tools": tool_names, **({"label": call_label} if call_label else {}), "content": result.get("content", ""), "tool_calls_made": result.get("tool_calls_made", []), diff --git a/backend/app/runner/llm.py b/backend/app/runner/llm.py index a52e25f..2fd2f43 100644 --- a/backend/app/runner/llm.py +++ b/backend/app/runner/llm.py @@ -34,6 +34,8 @@ def call_llm( tools: list[str] | None = None, on_event: Callable[[dict], None] | None = None, call_id: str | None = None, + tool_registry: dict | None = None, + tool_schemas_by_name: dict | None = None, **opts, ) -> dict: """ @@ -50,11 +52,23 @@ def call_llm( tagged with ``call_id``. call_id: unique id for this call, included on every emitted event so concurrent calls within one node can be disambiguated. + tool_registry: per-context callable registry; defaults to the process + runtime registry. Nodes pass a private copy containing their + embedded custom tools. + tool_schemas_by_name: schemas paired with ``tool_registry``; defaults + to the process runtime schemas. **opts: forwarded as additional fields in the request body. Returns: {content, messages, tool_calls_made, usage, cost} """ + registry = tool_registry if tool_registry is not None else REGISTRY + schemas_by_name = ( + tool_schemas_by_name + if tool_schemas_by_name is not None + else TOOL_SCHEMAS + ) + # Subscription-OAuth provider dispatch — the runner subprocess gets an # OAuth bearer + (for Codex) account id pre-resolved in env at spawn time. # Codex uses the Responses API, not chat completions, so the call shape @@ -69,8 +83,8 @@ def call_llm( model=model, prompt=prompt, tools=tools, - tool_registry=REGISTRY, - tool_schemas_by_name=TOOL_SCHEMAS, + tool_registry=registry, + tool_schemas_by_name=schemas_by_name, on_event=on_event, call_id=call_id, access_token=os.getenv("LLM_API_KEY", ""), @@ -102,7 +116,7 @@ def call_llm( messages = list(prompt) tools = tools or [] - tool_schemas = [TOOL_SCHEMAS[t] for t in tools if t in TOOL_SCHEMAS] + tool_schemas = [schemas_by_name[t] for t in tools if t in schemas_by_name] # Model limits drive compaction. Unknown model (catalog miss) → limits stay # zero and is_overflow() never fires, so a long node loop runs unchanged. @@ -249,7 +263,7 @@ def _emit(ev: dict) -> None: "round": round_idx, } ) - fn = REGISTRY.get(fn_name) + fn = registry.get(fn_name) if fn is None: # An MCP tool whose server didn't connect never lands in the # registry — report the server's state (and the needs_auth diff --git a/backend/app/runner/node_tools.py b/backend/app/runner/node_tools.py new file mode 100644 index 0000000..be49827 --- /dev/null +++ b/backend/app/runner/node_tools.py @@ -0,0 +1,67 @@ +"""Node-embedded tool declarations. + +Node source can define lowercase ``NodeTool`` subclasses alongside its +``run(inputs, ctx)`` function. Classes are registered as they are created +while that source is executed inside :func:`node_tool_registration`. + +The active registrar is a ContextVar rather than process-global mutable state: +independent nodes execute concurrently in worker threads and may legitimately +define tools with the same name. +""" +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Callable, Iterator + + +RegisterFn = Callable[[type["NodeTool"]], None] +_CURRENT_REGISTRAR: ContextVar[RegisterFn | None] = ContextVar( + "node_tool_registrar", default=None +) + + +class NodeTool: + """Base class for a custom tool embedded in a node's Python source. + + Subclasses use their lowercase class name as the provider-facing tool name + and declare ``description``, JSON-schema ``parameters``, and an + ``execute(self, ctx, **arguments)`` implementation. + """ + + description: str = "" + parameters: dict = {"type": "object", "properties": {}} + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + registrar = _CURRENT_REGISTRAR.get() + if registrar is not None: + registrar(cls) + + def execute(self, ctx, **arguments): # pragma: no cover - contract stub + raise NotImplementedError + + +@contextmanager +def node_tool_registration(register: RegisterFn) -> Iterator[None]: + """Route ``NodeTool`` subclasses created in this scope to ``register``.""" + + token = _CURRENT_REGISTRAR.set(register) + try: + yield + finally: + _CURRENT_REGISTRAR.reset(token) + + +def execute_node_source(code: str, register: RegisterFn) -> dict: + """Execute a node code blob and return its namespace. + + Keeping registration wrapped around the ``exec`` means class creation is + the declaration: no extra node field, ``NODE_TOOLS`` list, or namespace + scan is needed. + """ + + namespace: dict = {} + with node_tool_registration(register): + exec(code or "", namespace, namespace) + return namespace diff --git a/backend/tests/test_call_chats.py b/backend/tests/test_call_chats.py index 973239e..6409942 100644 --- a/backend/tests/test_call_chats.py +++ b/backend/tests/test_call_chats.py @@ -431,9 +431,10 @@ def test_send_turn_materializes_chat_lazily_and_persists(db_factory, monkeypatch captured = {} - def fake_start(turn_id, chat_id, messages, tools, model, child_env): + def fake_start(turn_id, chat_id, messages, tools, model, child_env, node_code=""): captured.update( - turn_id=turn_id, chat_id=chat_id, messages=messages, tools=tools, model=model + turn_id=turn_id, chat_id=chat_id, messages=messages, tools=tools, + model=model, node_code=node_code, ) monkeypatch.setattr(cc_api.chat_service, "start_chat_turn", fake_start) @@ -459,6 +460,7 @@ def fake_start(turn_id, chat_id, messages, tools, model, child_env): assert captured["model"] == "anthropic/claude-sonnet-4.5" assert captured["messages"][-1]["content"] == "now expand point 2" assert captured["tools"] == ["web_search"] + assert "def run(inputs, ctx)" in captured["node_code"] def test_send_turn_empty_message_400_without_materializing(db_factory): @@ -481,7 +483,7 @@ def test_send_turn_uses_switched_model_and_persists_selection(db_factory, monkey captured = {} - def fake_start(turn_id, chat_id, messages, tools, model, child_env): + def fake_start(turn_id, chat_id, messages, tools, model, child_env, node_code=""): captured.update(model=model) monkeypatch.setattr(cc_api.chat_service, "start_chat_turn", fake_start) @@ -595,7 +597,7 @@ def test_send_turn_coalesces_dangling_user_message(db_factory, monkeypatch): captured = {} - def fake_start(turn_id, chat_id, messages, tools, model, child_env): + def fake_start(turn_id, chat_id, messages, tools, model, child_env, node_code=""): captured.update(messages=messages) monkeypatch.setattr(cc_api.chat_service, "start_chat_turn", fake_start) diff --git a/backend/tests/test_node_tools.py b/backend/tests/test_node_tools.py new file mode 100644 index 0000000..a6d83c9 --- /dev/null +++ b/backend/tests/test_node_tools.py @@ -0,0 +1,209 @@ +"""Node-embedded custom tool registration and dispatch.""" +from __future__ import annotations + +import pytest + +from app.runner import llm as llm_mod +from app.runner import chat_child +from app.runner.ctx import Ctx +from app.runner.node_tools import execute_node_source +from app.runner.tools import REGISTRY + + +TOOL_SOURCE = ''' +from app.runner.node_tools import NodeTool + +class lookup_order(NodeTool): + description = "Look up an order by id." + parameters = { + "type": "object", + "properties": {"order_id": {"type": "string"}}, + "required": ["order_id"], + } + + def execute(self, ctx, order_id): + ctx.log(f"looked up {order_id}") + return {"order_id": order_id, "status": "shipped"} + +def run(inputs, ctx): + result = ctx.agent( + prompt=f"Help with order {inputs['order_id']}", + tools=["lookup_order"], + ) + return {"answer": result["content"]} +''' + + +def test_embedded_tool_auto_registers_on_source_execution(tmp_path): + ctx = Ctx(workdir=tmp_path, default_model="test-model") + + namespace = execute_node_source(TOOL_SOURCE, ctx.register_node_tool) + + assert "lookup_order" in ctx.tool_registry + assert "lookup_order" in ctx.tool_schemas + assert ctx.tool_schemas["lookup_order"]["function"] == { + "name": "lookup_order", + "description": "Look up an order by id.", + "parameters": { + "type": "object", + "properties": {"order_id": {"type": "string"}}, + "required": ["order_id"], + }, + } + assert callable(namespace["run"]) + # Registration is local to this node context, never process-global. + assert "lookup_order" not in REGISTRY + + +def test_custom_tool_uses_existing_quoted_agent_surface(tmp_path, monkeypatch): + ctx = Ctx(workdir=tmp_path, default_model="test-model") + namespace = execute_node_source(TOOL_SOURCE, ctx.register_node_tool) + captured = {} + + def fake_call_llm( + model, + prompt, + tools, + tool_registry, + tool_schemas_by_name, + **kwargs, + ): + captured.update( + model=model, + prompt=prompt, + tools=tools, + registry=tool_registry, + schemas=tool_schemas_by_name, + ) + tool_result = tool_registry["lookup_order"](order_id="ord-7") + return { + "content": tool_result["status"], + "messages": [ + {"role": "user", "content": prompt}, + {"role": "assistant", "content": tool_result["status"]}, + ], + "tool_calls_made": [{ + "name": "lookup_order", + "args": {"order_id": "ord-7"}, + "result": tool_result, + }], + "usage": {}, + "cost": 0.0, + } + + monkeypatch.setattr(llm_mod, "call_llm", fake_call_llm) + + result = namespace["run"]({"order_id": "ord-7"}, ctx) + + assert result == {"answer": "shipped"} + assert captured["tools"] == ["lookup_order"] + assert captured["registry"] is ctx.tool_registry + assert captured["schemas"] is ctx.tool_schemas + assert ctx.llm_calls[0]["tools"] == ["lookup_order"] + assert ctx.tool_calls[0]["name"] == "lookup_order" + assert "looked up ord-7" in ctx.logs + + +def test_custom_tool_is_available_through_direct_proxy(tmp_path): + ctx = Ctx(workdir=tmp_path, default_model="test-model") + execute_node_source(TOOL_SOURCE, ctx.register_node_tool) + + # Direct custom calls follow the native tools' positional-or-keyword + # behavior even though agentic dispatch always supplies JSON kwargs. + result = ctx.tools.lookup_order("ord-9") + + assert result == {"order_id": "ord-9", "status": "shipped"} + assert ctx.tool_calls[-1] == { + "name": "lookup_order", + "args": {"order_id": "ord-9"}, + "via": "direct", + "result": result, + } + + +def test_same_custom_name_is_isolated_between_node_contexts(tmp_path): + source = ''' +from app.runner.node_tools import NodeTool +class local_tool(NodeTool): + description = "Return this node's value." + parameters = {"type": "object", "properties": {}} + def execute(self, ctx): + return {"value": VALUE} +''' + first = Ctx(workdir=tmp_path, default_model="test-model") + second = Ctx(workdir=tmp_path, default_model="test-model") + + execute_node_source("VALUE = 'first'\n" + source, first.register_node_tool) + execute_node_source("VALUE = 'second'\n" + source, second.register_node_tool) + + assert first.tools.local_tool() == {"value": "first"} + assert second.tools.local_tool() == {"value": "second"} + assert first.tool_registry["local_tool"] is not second.tool_registry["local_tool"] + + +def test_continued_chat_restores_tool_from_frozen_node_source(tmp_path, monkeypatch): + events = [] + seen = {} + messages = [{"role": "user", "content": "continue"}] + + monkeypatch.setattr(chat_child, "_install_sigterm_handler", lambda: None) + monkeypatch.setattr(chat_child, "_load_mcp_tools", lambda: None) + monkeypatch.setattr(chat_child, "_emit", events.append) + monkeypatch.setattr(chat_child, "_read_payload", lambda: { + "messages": messages, + "tools": ["lookup_order"], + "node_code": TOOL_SOURCE, + "model": "test-model", + "workdir": str(tmp_path), + }) + + def fake_agent(self, model=None, prompt=None, tools=None, **kwargs): + seen["tools"] = tools + seen["result"] = self.tool_registry["lookup_order"](order_id="ord-11") + return { + "messages": [*messages, {"role": "assistant", "content": "done"}], + "usage": {}, + "cost": 0.0, + } + + monkeypatch.setattr(Ctx, "agent", fake_agent) + + chat_child.main() + + assert seen == { + "tools": ["lookup_order"], + "result": {"order_id": "ord-11", "status": "shipped"}, + } + assert events[-1]["type"] == "run_finished" + assert events[-1]["status"] == "success" + + +@pytest.mark.parametrize("name", ["LookupOrder", "lookupOrder", "LOOKUP_ORDER"]) +def test_custom_tool_requires_native_lowercase_casing(tmp_path, name): + source = f''' +from app.runner.node_tools import NodeTool +class {name}(NodeTool): + description = "bad casing" + parameters = {{"type": "object", "properties": {{}}}} + def execute(self, ctx): + return {{}} +''' + ctx = Ctx(workdir=tmp_path, default_model="test-model") + + with pytest.raises(ValueError, match="lowercase snake_case"): + execute_node_source(source, ctx.register_node_tool) + + +def test_custom_tool_cannot_shadow_native_tool(tmp_path): + source = ''' +from app.runner.node_tools import NodeTool +class shell(NodeTool): + description = "shadow shell" + parameters = {"type": "object", "properties": {}} + def execute(self, ctx): + return {} +''' + ctx = Ctx(workdir=tmp_path, default_model="test-model") + + with pytest.raises(ValueError, match="conflicts with an existing tool"): + execute_node_source(source, ctx.register_node_tool) diff --git a/backend/tests/test_orchestrator.py b/backend/tests/test_orchestrator.py index b2a7400..f5c4f82 100644 --- a/backend/tests/test_orchestrator.py +++ b/backend/tests/test_orchestrator.py @@ -1891,6 +1891,15 @@ def test_system_prompt_offers_both_direct_and_llm_tool_forms(): assert "tools=[...]" in p or "tools=[" in p +def test_system_prompt_teaches_embedded_custom_tool_contract(): + p = SYSTEM_PROMPT + assert "from app.runner.node_tools import NodeTool" in p + assert "class lookup_order(NodeTool)" in p + assert 'tools=["lookup_order"]' in p + assert "lowercase snake_case" in p + assert "registers it automatically" in p + + def test_system_prompt_lists_node_runtime_tool_signatures(): """The orchestrator writes Python like `ctx.tools.web_fetch(...)` — it needs the canonical signatures (param names + types) in the prompt or it diff --git a/backend/tests/test_runner.py b/backend/tests/test_runner.py index af4c6d1..d3b7633 100644 --- a/backend/tests/test_runner.py +++ b/backend/tests/test_runner.py @@ -246,6 +246,45 @@ def run(inputs, ctx): assert any(tc["name"] == "shell" for tc in nr["tool_calls"]) +def test_node_embedded_tool_runs_in_child_subprocess(): + code = ''' +from app.runner.node_tools import NodeTool + +class lookup_order(NodeTool): + description = "Look up an order by id." + parameters = { + "type": "object", + "properties": {"order_id": {"type": "string"}}, + "required": ["order_id"], + } + + def execute(self, ctx, order_id): + return {"status": f"{order_id}:shipped"} + +def run(inputs, ctx): + result = ctx.tools.lookup_order(order_id=inputs["order_id"]) + return {"status": result["status"]} +''' + wf = { + "id": "wf", + "input_node_id": "a", + "output_node_id": "a", + "nodes": [make_node( + "a", + code, + inputs=[{"name": "order_id", "required": True}], + outputs=[{"name": "status"}], + )], + "edges": [], + } + + result = run_workflow_sync(wf, {"order_id": "ord-3"}) + + assert result["status"] == "success", result + assert result["outputs"] == {"status": "ord-3:shipped"} + assert result["node_runs"][0]["tool_calls"][0]["name"] == "lookup_order" + + def test_run_workdir_is_removed(tmp_path, monkeypatch): workdir = tmp_path / "wfrun-test" monkeypatch.setattr(runner_mod.tempfile, "mkdtemp", lambda prefix: str(workdir))