From 341055469d3beb9f5d3aade57dcc7a4c49c9e593 Mon Sep 17 00:00:00 2001 From: Akshat malik Date: Sun, 13 Sep 2026 12:33:27 +0530 Subject: [PATCH 1/3] fix: correctly bind async local plugins as coroutines --- .../engine/langgraph/graph/graph_builder.py | 5 ++- tests/test_review_regression.py | 39 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) create mode 100644 tests/test_review_regression.py diff --git a/src/agent_engine/engine/langgraph/graph/graph_builder.py b/src/agent_engine/engine/langgraph/graph/graph_builder.py index 81187057..f1f34b24 100644 --- a/src/agent_engine/engine/langgraph/graph/graph_builder.py +++ b/src/agent_engine/engine/langgraph/graph/graph_builder.py @@ -255,7 +255,10 @@ def _build_agent_tools( server_by_tool: dict[str, str] = {} for tool_spec in spec.tools: function = self._tool_loader.load(tool_spec.id) - tools.append(StructuredTool.from_function(function, description=tool_spec.description)) + if inspect.iscoroutinefunction(function): + tools.append(StructuredTool.from_function(coroutine=function, description=tool_spec.description)) + else: + tools.append(StructuredTool.from_function(func=function, description=tool_spec.description)) for mcp in spec.mcps: server_tools = self._mcp_tools.get(mcp.id, []) tools.extend(server_tools) diff --git a/tests/test_review_regression.py b/tests/test_review_regression.py new file mode 100644 index 00000000..7261467a --- /dev/null +++ b/tests/test_review_regression.py @@ -0,0 +1,39 @@ +import pytest +from langchain_core.messages import AIMessage, ToolMessage + +from agent_engine.engine.langgraph.engine import LangGraphEngine +from tests.approvals.test_engine_hitl import _spec as approval_spec + + +class SingleToolCallModel: + def __init__(self): + self.calls = 0 + + def bind_tools(self, tools): + return self + + async def ainvoke(self, messages): + if any(isinstance(m, ToolMessage) for m in messages): + return AIMessage(content='done') + self.calls += 1 + return AIMessage(content='', tool_calls=[{ + 'name': 'record_value', 'args': {'message': 'reviewed-value'}, + 'id': f'call-{self.calls}', 'type': 'tool_call', + }]) + + +@pytest.mark.asyncio +async def test_async_local_tool_executes(tmp_path): + output = tmp_path / 'async-executed.txt' + tool = tmp_path / 'plugins/tools/record_value.py' + tool.parent.mkdir(parents=True) + tool.write_text('from pathlib import Path\n' + 'async def record_value(message: str) -> str:\n' + f' Path({str(output)!r}).write_text(message)\n' + ' return message\n') + model = SingleToolCallModel() + async with LangGraphEngine(tmp_path, model_factory=lambda *a, **kw: model) as engine: + await engine.build(approval_spec('record_value', auto_mode=True)) + result = await engine.run('record a value') + print(f'async tool executed={output.exists()}, used_tools={result.used_tools}') + assert output.exists() From f078016147d9b13088a1d7b922ba5abdb85bf23d Mon Sep 17 00:00:00 2001 From: Akshat malik Date: Mon, 5 Oct 2026 12:39:33 +0530 Subject: [PATCH 2/3] added more regression tests --- tests/engine/test_async_local_tools.py | 228 +++++++++++++++++++++++++ tests/test_review_regression.py | 39 ----- 2 files changed, 228 insertions(+), 39 deletions(-) create mode 100644 tests/engine/test_async_local_tools.py delete mode 100644 tests/test_review_regression.py diff --git a/tests/engine/test_async_local_tools.py b/tests/engine/test_async_local_tools.py new file mode 100644 index 00000000..304f21bd --- /dev/null +++ b/tests/engine/test_async_local_tools.py @@ -0,0 +1,228 @@ +"""Regression tests for async and sync local tool execution in the runtime engine.""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Any, cast + +import pytest +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import AIMessage, ToolMessage +from langchain_core.messages.tool import ToolCall as LCToolCall + +from agent_engine.core.spec import ( + AgentSpec, + BasePromptSet, + GraphNode, + ModelConfig, + SystemMeta, + SystemSpec, + ToolSpec, +) +from agent_engine.engine.langgraph.engine import LangGraphEngine + +_MODEL = ModelConfig(provider="fake", name="fake", temperature=None) + + +class SingleToolCallModel: + """Fake chat model that invokes a specified tool once and inspects returned messages.""" + + def __init__( + self, + tool_name: str, + tool_args: dict[str, Any] | None = None, + ) -> None: + self.tool_name = tool_name + self.tool_args = tool_args or {"message": "reviewed-value"} + self.calls = 0 + self.received_tool_messages: list[ToolMessage] = [] + + def bind_tools(self, tools: list[Any]) -> SingleToolCallModel: + return self + + async def ainvoke(self, messages: list[Any]) -> AIMessage: + for m in messages: + if isinstance(m, ToolMessage): + self.received_tool_messages.append(m) + return AIMessage(content=f"model-received: {m.content}") + + self.calls += 1 + return AIMessage( + content="", + tool_calls=[ + LCToolCall( + name=self.tool_name, + args=self.tool_args, + id=f"call-{self.calls}", + ) + ], + ) + + +def _spec(tool_id: str, *, auto_mode: bool = True) -> SystemSpec: + agent = AgentSpec( + id="tool_agent", + name="tool_agent", + description="test tool agent", + model=_MODEL, + prompts=BasePromptSet(), + tools=(ToolSpec(tool_id, f"Tool {tool_id}"),), + auto_mode=auto_mode, + ) + return SystemSpec( + meta=SystemMeta(name="async_tool_test"), + defaults=None, + graph=GraphNode(node=agent), + ) + + +def _write_tool(base_dir: Path, tool_id: str, code: str) -> None: + tools_dir = base_dir / "plugins" / "tools" + tools_dir.mkdir(parents=True, exist_ok=True) + (tools_dir / f"{tool_id}.py").write_text(code, encoding="utf-8") + + +@pytest.mark.asyncio +async def test_async_local_tool_executes_once_and_returns_value_to_model( + tmp_path: Path, +) -> None: + counter_file = tmp_path / "calls.txt" + _write_tool( + tmp_path, + "record_value", + "from pathlib import Path\n" + "async def record_value(message: str) -> str:\n" + f" with open({str(counter_file)!r}, 'a') as f:\n" + " f.write('x')\n" + " return f'processed-{message}'\n", + ) + model = SingleToolCallModel("record_value", {"message": "reviewed-value"}) + async with LangGraphEngine( + tmp_path, model_factory=lambda *a, **kw: cast(BaseChatModel, model) + ) as engine: + await engine.build(_spec("record_value")) + result = await engine.run("run tool") + + # 1. The async tool executes exactly once + assert counter_file.exists() + assert counter_file.read_text() == "x" + + # 2. The resolved return value reaches the model as ToolMessage content (not coroutine object) + assert len(model.received_tool_messages) == 1 + tool_msg = model.received_tool_messages[0] + assert tool_msg.content == "processed-reviewed-value" + assert "coroutine" not in str(tool_msg.content).lower() + + # Model answer reflects tool output and tool usage status is succeeded + assert result.answer == "model-received: processed-reviewed-value" + assert len(result.used_tools) == 1 + assert result.used_tools[0].status == "succeeded" + + +@pytest.mark.asyncio +async def test_async_local_tool_exception_recorded_as_failed( + tmp_path: Path, +) -> None: + _write_tool( + tmp_path, + "failing_tool", + "async def failing_tool(message: str) -> str:\n" + " raise ValueError(f'async failure: {message}')\n", + ) + model = SingleToolCallModel("failing_tool", {"message": "bad-arg"}) + async with LangGraphEngine( + tmp_path, model_factory=lambda *a, **kw: cast(BaseChatModel, model) + ) as engine: + await engine.build(_spec("failing_tool")) + result = await engine.run("run failing tool") + + # 3. An exception raised inside an async local tool is recorded as a failed tool execution + assert len(model.received_tool_messages) == 1 + tool_msg = model.received_tool_messages[0] + assert "Tool error: async failure: bad-arg" in str(tool_msg.content) + + assert len(result.used_tools) == 1 + assert result.used_tools[0].status == "failed" + assert "async failure: bad-arg" in (result.used_tools[0].error or "") + + +@pytest.mark.asyncio +async def test_async_local_tool_cancellation_propagates( + tmp_path: Path, +) -> None: + started_file = tmp_path / "started.txt" + cancelled_file = tmp_path / "cancelled.txt" + + _write_tool( + tmp_path, + "cancelling_tool", + "import asyncio\n" + "from pathlib import Path\n" + "async def cancelling_tool(message: str) -> str:\n" + f" Path({str(started_file)!r}).write_text('started')\n" + " try:\n" + " await asyncio.sleep(10)\n" + " except asyncio.CancelledError:\n" + f" Path({str(cancelled_file)!r}).write_text('cancelled')\n" + " raise\n" + " return 'done'\n", + ) + + model = SingleToolCallModel("cancelling_tool", {"message": "test"}) + async with LangGraphEngine( + tmp_path, model_factory=lambda *a, **kw: cast(BaseChatModel, model) + ) as engine: + await engine.build(_spec("cancelling_tool")) + task = asyncio.create_task(engine.run("run tool")) + + # Wait until the tool actually starts execution + for _ in range(50): + if started_file.exists(): + break + await asyncio.sleep(0.05) + + assert started_file.exists() + + # Cancel the execution task + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # 4. Cancellation propagates to the async local tool + assert cancelled_file.exists() + assert cancelled_file.read_text() == "cancelled" + + +@pytest.mark.asyncio +async def test_sync_local_tool_follows_existing_execution_path( + tmp_path: Path, +) -> None: + counter_file = tmp_path / "sync_calls.txt" + _write_tool( + tmp_path, + "sync_tool", + "from pathlib import Path\n" + "def sync_tool(message: str) -> str:\n" + f" with open({str(counter_file)!r}, 'a') as f:\n" + " f.write('s')\n" + " return f'sync-{message}'\n", + ) + model = SingleToolCallModel("sync_tool", {"message": "reviewed-value"}) + async with LangGraphEngine( + tmp_path, model_factory=lambda *a, **kw: cast(BaseChatModel, model) + ) as engine: + await engine.build(_spec("sync_tool")) + result = await engine.run("run sync tool") + + # 5. Synchronous local tools still follow existing execution path + assert counter_file.exists() + assert counter_file.read_text() == "s" + + assert len(model.received_tool_messages) == 1 + tool_msg = model.received_tool_messages[0] + assert tool_msg.content == "sync-reviewed-value" + + assert result.answer == "model-received: sync-reviewed-value" + assert len(result.used_tools) == 1 + assert result.used_tools[0].status == "succeeded" diff --git a/tests/test_review_regression.py b/tests/test_review_regression.py deleted file mode 100644 index 7261467a..00000000 --- a/tests/test_review_regression.py +++ /dev/null @@ -1,39 +0,0 @@ -import pytest -from langchain_core.messages import AIMessage, ToolMessage - -from agent_engine.engine.langgraph.engine import LangGraphEngine -from tests.approvals.test_engine_hitl import _spec as approval_spec - - -class SingleToolCallModel: - def __init__(self): - self.calls = 0 - - def bind_tools(self, tools): - return self - - async def ainvoke(self, messages): - if any(isinstance(m, ToolMessage) for m in messages): - return AIMessage(content='done') - self.calls += 1 - return AIMessage(content='', tool_calls=[{ - 'name': 'record_value', 'args': {'message': 'reviewed-value'}, - 'id': f'call-{self.calls}', 'type': 'tool_call', - }]) - - -@pytest.mark.asyncio -async def test_async_local_tool_executes(tmp_path): - output = tmp_path / 'async-executed.txt' - tool = tmp_path / 'plugins/tools/record_value.py' - tool.parent.mkdir(parents=True) - tool.write_text('from pathlib import Path\n' - 'async def record_value(message: str) -> str:\n' - f' Path({str(output)!r}).write_text(message)\n' - ' return message\n') - model = SingleToolCallModel() - async with LangGraphEngine(tmp_path, model_factory=lambda *a, **kw: model) as engine: - await engine.build(approval_spec('record_value', auto_mode=True)) - result = await engine.run('record a value') - print(f'async tool executed={output.exists()}, used_tools={result.used_tools}') - assert output.exists() From 02ef9ee6ad358b0946640789a0b61508287e31bd Mon Sep 17 00:00:00 2001 From: Akshat malik Date: Mon, 5 Oct 2026 23:11:38 +0530 Subject: [PATCH 3/3] style: format graph_builder.py with ruff --- .../engine/langgraph/graph/graph_builder.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/agent_engine/engine/langgraph/graph/graph_builder.py b/src/agent_engine/engine/langgraph/graph/graph_builder.py index f1f34b24..78bef498 100644 --- a/src/agent_engine/engine/langgraph/graph/graph_builder.py +++ b/src/agent_engine/engine/langgraph/graph/graph_builder.py @@ -256,9 +256,15 @@ def _build_agent_tools( for tool_spec in spec.tools: function = self._tool_loader.load(tool_spec.id) if inspect.iscoroutinefunction(function): - tools.append(StructuredTool.from_function(coroutine=function, description=tool_spec.description)) + tools.append( + StructuredTool.from_function( + coroutine=function, description=tool_spec.description + ) + ) else: - tools.append(StructuredTool.from_function(func=function, description=tool_spec.description)) + tools.append( + StructuredTool.from_function(func=function, description=tool_spec.description) + ) for mcp in spec.mcps: server_tools = self._mcp_tools.get(mcp.id, []) tools.extend(server_tools)