Enforce async-only function registration in FunctionRegistry#373
Conversation
register_function now raises ValueError when given a synchronous function. call_function and _run_one_tool are simplified since all registered callables are guaranteed to be coroutines.
📝 WalkthroughWalkthroughThe pull request migrates the function calling system from synchronous to asynchronous execution. Core registry and LLM classes now enforce async function registration and provide async call methods. Test files and examples throughout the codebase are updated to register and invoke async functions with await syntax. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
agents-core/vision_agents/testing/_mock_tools.py (2)
27-66:⚠️ Potential issue | 🟠 Major
mock_toolssilently fails when given a synchronous callable.
mock_toolsassignsmock_fndirectly tofunc_def.function(Line 60) without any async check. Every execution path —_run_one_tool(await fn(**args)) andFunctionRegistry.call_function(await func_def.function(...)) — unconditionally awaits the stored callable. Passing a sync callable produces aTypeError(e.g.,"object str can't be used in 'await' expression") that is caught by_run_one_tool's broadexceptand silently returned as{"error": ...}, making test assertions fail in a confusing way.
mock_functionsis safe because it wraps everything inAsyncMock(side_effect=...).mock_toolshas no equivalent guard.🛡️ Proposed fix — mirror the registry's own enforcement
+import inspect from collections.abc import Callable, Generator +from collections.abc import Coroutine ... `@contextmanager` def mock_tools( llm: LLM, - mocks: dict[str, Callable[..., Any]], + mocks: dict[str, Callable[..., Coroutine[Any, Any, Any]]], ) -> Generator[None, None, None]: ... for tool_name in mocks: if registry._functions.get(tool_name) is None: raise KeyError(f"Tool '{tool_name}' is not registered on this LLM") + + for tool_name, mock_fn in mocks.items(): + if not inspect.iscoroutinefunction(mock_fn): + raise ValueError( + f"Mock for '{tool_name}' must be an async function" + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@agents-core/vision_agents/testing/_mock_tools.py` around lines 27 - 66, mock_tools assigns the provided mock_fn directly to func_def.function (in mock_tools) without ensuring it's async, so callers like _run_one_tool and FunctionRegistry.call_function that use await will raise a TypeError for synchronous callables; fix by wrapping non-async mocks in an async wrapper (or convert them to AsyncMock) before assignment so func_def.function is always awaitable—detect with asyncio.iscoroutinefunction or inspect.iscoroutinefunction and if false, replace mock_fn with an async wrapper that calls the sync mock (or use AsyncMock(side_effect=mock_fn)) before setting registry._functions[tool_name].function, mirroring the registry's enforcement used by mock_functions.
27-66:⚠️ Potential issue | 🟠 Major
mock_toolssilently breaks when given a synchronous callable.
mock_toolsassigns the raw callable directly tofunc_def.function(line 60). Every downstream execution path (_run_one_tool,FunctionRegistry.call_function) now unconditionallyawaits that callable. Passing a sync function produces aTypeErrorat runtime (e.g.,"object str can't be used in 'await' expression"), which is silently swallowed by_run_one_tool's broadexceptclause and surfaces as an unexpected{"error": ...}result in tests.
mock_functionsis safe because it wraps withAsyncMock(side_effect=...), which handles sync callables;mock_toolshas no such guard.Consider adding a validation check (mirroring the registry's own enforcement) or updating the type hint to communicate the constraint clearly:
🛡️ Proposed fix
+import inspect +from collections.abc import Coroutine from collections.abc import Callable, Generator ... `@contextmanager` def mock_tools( llm: LLM, - mocks: dict[str, Callable[..., Any]], + mocks: dict[str, Callable[..., Coroutine[Any, Any, Any]]], ) -> Generator[None, None, None]: ... for tool_name in mocks: if registry._functions.get(tool_name) is None: raise KeyError(f"Tool '{tool_name}' is not registered on this LLM") + for tool_name, mock_fn in mocks.items(): + if not inspect.iscoroutinefunction(mock_fn): + raise ValueError( + f"Mock for '{tool_name}' must be an async function" + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@agents-core/vision_agents/testing/_mock_tools.py` around lines 27 - 66, mock_tools assigns raw callables to func_def.function but downstream code always awaits them, so passing a synchronous function causes a TypeError; fix by detecting non-coroutine mocks in mock_tools (use inspect.iscoroutinefunction or asyncio.iscoroutinefunction) and either wrap sync callables with AsyncMock(side_effect=mock_fn) before assigning to func_def.function or raise a clear TypeError; update the behavior in mock_tools where it assigns to registry._functions[tool_name].function (and adjust the typing/comments if you prefer the stricter error approach) so tests no longer get swallowed await errors.tests/test_testing/test_mock_tools.py (1)
25-45:⚠️ Potential issue | 🟠 MajorTest design still depends on mocking despite no-mock test policy.
These updated tests continue to validate behavior via mocked callables/
AsyncMock, which violates the repository’s testing rule and should be replaced with concrete pytest fixtures/fakes.As per coding guidelines, "Never mock in tests; use pytest for testing".
Also applies to: 127-188
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_testing/test_mock_tools.py` around lines 25 - 45, The test test_mock_tools_swaps_and_restores uses a mocked callable (mock_fn/AsyncMock) via mock_tools which violates the no-mock rule; replace the mock with a concrete fake implementation and/or pytest fixture: register a real async function (e.g., a simple async def replacement_tool(x: int) -> int that returns x * 100) and use that with mock_tools(llm, {"my_tool": replacement_tool}); ensure you still capture original_fn from llm.function_registry._functions["my_tool"].function, assert the active function is the replacement function, call the function via llm.function_registry.call_function("my_tool", {"x": 5}) to verify behavior, and after the context verify restoration to original_fn—apply the same replacement approach to the other tests noted (lines ~127-188) that currently use AsyncMock or other mocks.tests/test_function_calling.py (1)
177-370:⚠️ Potential issue | 🟠 MajorReplace
Mock/patch-based tests with pytest-native non-mocking fixtures.These changed integration-style tests still rely on
unittest.mock(@patch,Mock), which conflicts with the repository test policy and should be refactored to non-mocking test doubles/fixtures.As per coding guidelines, "Never mock in tests; use pytest for testing".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_function_calling.py` around lines 177 - 370, The tests (test_openai_function_calling_response, test_openai_conversational_response, test_claude_function_calling_response, test_claude_conversational_response, test_gemini_function_calling_response) currently use unittest.mock patching and Mock objects; refactor them to use pytest-native fixtures/monkeypatch instead: remove `@patch` decorators and Mock usage, add pytest fixtures (e.g., openai_client_fixture, anthropic_client_fixture, gemini_client_fixture) that create lightweight fake client objects implementing the minimal methods used by OpenAILLM.responses.create, ClaudeLLM.messages.create, and GeminiLLM.aio.chats.create/send_message_stream (including returning an async iterator for Gemini), and inject those fixtures into tests via function arguments or monkeypatch the module attributes (e.g., set vision_agents.plugins.openai.openai_llm.AsyncOpenAI to the fixture or patch the client attribute on the LLM instance) so behavior is deterministic without unittest.mock.
🧹 Nitpick comments (2)
agents-core/vision_agents/core/llm/llm.py (2)
247-249:_run_one_toolbypasses required-parameter validation.
get_callablereturns the raw function without the required-parameter check thatFunctionRegistry.call_functionperforms for introspectedFunctionDefinitionentries. When a required parameter is absent, Python raises a genericTypeErrorfrom theawait fn(**args)call rather than the descriptive message produced by the registry's own guard. The error is swallowed by the broadexceptand returned as{"error": ...}, so this is not a crash — but the inconsistency between_run_one_tooland the publicLLM.call_functionis confusing.Swapping to
self.function_registry.call_functionunifies both paths and makes the behavior predictable:♻️ Proposed fix
async def _invoke(): - fn = self.function_registry.get_callable(tc["name"]) - return await fn(**args) + return await self.function_registry.call_function(tc["name"], args)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@agents-core/vision_agents/core/llm/llm.py` around lines 247 - 249, _run_one_tool currently uses self.function_registry.get_callable(tc["name"]) and awaits the raw fn(**args), which bypasses the FunctionRegistry parameter validation provided by FunctionRegistry.call_function and leads to non-descriptive TypeError on missing required params; change _run_one_tool to call and await self.function_registry.call_function(tc["name"], args) (or the exact call_function signature used elsewhere) instead of get_callable to ensure FunctionDefinition validation runs and errors match LLM.call_function behavior, preserving the existing exception handling and return shape.
247-249:_run_one_toolbypasses required-parameter validation.Using
get_callableand thenawait fn(**args)directly skips the required-parameter guard inFunctionRegistry.call_function, which iteratesfunc_def.parametersand raises a descriptiveTypeErrorfor missing required args. When a required parameter is absent here, Python raises a genericTypeErrorthat is swallowed by the outerexceptand returned as{"error": ...}. The publicLLM.call_functionpath correctly delegates toFunctionRegistry.call_functionand gets the validation. The two paths are now inconsistent.♻️ Proposed fix — unify through the validated path
async def _invoke(): - fn = self.function_registry.get_callable(tc["name"]) - return await fn(**args) + return await self.function_registry.call_function(tc["name"], args)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@agents-core/vision_agents/core/llm/llm.py` around lines 247 - 249, The helper _run_one_tool currently fetches a raw callable via function_registry.get_callable and invokes it directly, bypassing required-parameter validation; change it to delegate to FunctionRegistry.call_function (e.g., call self.function_registry.call_function with the function name tc["name"] and the args dict) so the registry's parameter checks run, preserving the existing async/await and outer error handling; ensure you stop using get_callable in _run_one_tool and pass the same args through the validated call path used by LLM.call_function.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@agents-core/vision_agents/testing/_mock_tools.py`:
- Around line 27-66: mock_tools assigns the provided mock_fn directly to
func_def.function (in mock_tools) without ensuring it's async, so callers like
_run_one_tool and FunctionRegistry.call_function that use await will raise a
TypeError for synchronous callables; fix by wrapping non-async mocks in an async
wrapper (or convert them to AsyncMock) before assignment so func_def.function is
always awaitable—detect with asyncio.iscoroutinefunction or
inspect.iscoroutinefunction and if false, replace mock_fn with an async wrapper
that calls the sync mock (or use AsyncMock(side_effect=mock_fn)) before setting
registry._functions[tool_name].function, mirroring the registry's enforcement
used by mock_functions.
- Around line 27-66: mock_tools assigns raw callables to func_def.function but
downstream code always awaits them, so passing a synchronous function causes a
TypeError; fix by detecting non-coroutine mocks in mock_tools (use
inspect.iscoroutinefunction or asyncio.iscoroutinefunction) and either wrap sync
callables with AsyncMock(side_effect=mock_fn) before assigning to
func_def.function or raise a clear TypeError; update the behavior in mock_tools
where it assigns to registry._functions[tool_name].function (and adjust the
typing/comments if you prefer the stricter error approach) so tests no longer
get swallowed await errors.
In `@tests/test_function_calling.py`:
- Around line 177-370: The tests (test_openai_function_calling_response,
test_openai_conversational_response, test_claude_function_calling_response,
test_claude_conversational_response, test_gemini_function_calling_response)
currently use unittest.mock patching and Mock objects; refactor them to use
pytest-native fixtures/monkeypatch instead: remove `@patch` decorators and Mock
usage, add pytest fixtures (e.g., openai_client_fixture,
anthropic_client_fixture, gemini_client_fixture) that create lightweight fake
client objects implementing the minimal methods used by
OpenAILLM.responses.create, ClaudeLLM.messages.create, and
GeminiLLM.aio.chats.create/send_message_stream (including returning an async
iterator for Gemini), and inject those fixtures into tests via function
arguments or monkeypatch the module attributes (e.g., set
vision_agents.plugins.openai.openai_llm.AsyncOpenAI to the fixture or patch the
client attribute on the LLM instance) so behavior is deterministic without
unittest.mock.
In `@tests/test_testing/test_mock_tools.py`:
- Around line 25-45: The test test_mock_tools_swaps_and_restores uses a mocked
callable (mock_fn/AsyncMock) via mock_tools which violates the no-mock rule;
replace the mock with a concrete fake implementation and/or pytest fixture:
register a real async function (e.g., a simple async def replacement_tool(x:
int) -> int that returns x * 100) and use that with mock_tools(llm, {"my_tool":
replacement_tool}); ensure you still capture original_fn from
llm.function_registry._functions["my_tool"].function, assert the active function
is the replacement function, call the function via
llm.function_registry.call_function("my_tool", {"x": 5}) to verify behavior, and
after the context verify restoration to original_fn—apply the same replacement
approach to the other tests noted (lines ~127-188) that currently use AsyncMock
or other mocks.
---
Nitpick comments:
In `@agents-core/vision_agents/core/llm/llm.py`:
- Around line 247-249: _run_one_tool currently uses
self.function_registry.get_callable(tc["name"]) and awaits the raw fn(**args),
which bypasses the FunctionRegistry parameter validation provided by
FunctionRegistry.call_function and leads to non-descriptive TypeError on missing
required params; change _run_one_tool to call and await
self.function_registry.call_function(tc["name"], args) (or the exact
call_function signature used elsewhere) instead of get_callable to ensure
FunctionDefinition validation runs and errors match LLM.call_function behavior,
preserving the existing exception handling and return shape.
- Around line 247-249: The helper _run_one_tool currently fetches a raw callable
via function_registry.get_callable and invokes it directly, bypassing
required-parameter validation; change it to delegate to
FunctionRegistry.call_function (e.g., call self.function_registry.call_function
with the function name tc["name"] and the args dict) so the registry's parameter
checks run, preserving the existing async/await and outer error handling; ensure
you stop using get_callable in _run_one_tool and pass the same args through the
validated call path used by LLM.call_function.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Disabled knowledge base sources:
- Linear integration is disabled
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (13)
agents-core/vision_agents/core/llm/function_registry.pyagents-core/vision_agents/core/llm/llm.pyagents-core/vision_agents/testing/_mock_tools.pyplugins/aws/example/aws_realtime_function_calling_example.pyplugins/aws/tests/test_aws_realtime.pyplugins/gemini/tests/test_realtime_function_calling.pyplugins/openrouter/example/openrouter_example.pyplugins/openrouter/tests/test_openrouter_llm.pyplugins/xai/tests/test_xai_realtime.pyplugins/xai/tests/test_xai_tools.pytests/test_function_calling.pytests/test_openai_function_calling_integration.pytests/test_testing/test_mock_tools.py
2 participants
Motivation
register_functionpreviously accepted both sync and async callables. This led to hidden complexity:_run_one_toolhad to inspect every function at call time (iscoroutinefunction), branch intoawaitorasyncio.to_thread, and maintain a_maybe_awaithelper for the fallback path. All of this existed solely to paper over sync functions being registered in an async-first framework.By requiring async functions at registration time we:
ValueErrorat registration is easier to debug than unexpected blocking at runtime._invoke,_run_one_tool, andcall_functionno longer need sync/async branching;_maybe_awaitis removed entirely.asyncio.to_threadthemselves inside their async function, which makes the cost visible at the call site.Migration
Changes
FunctionRegistry.register— raisesValueErrorfor non-coroutine callables.FunctionRegistry.call_function— nowasync, awaits the registered function directly.LLM.call_function— nowasync.LLM._run_one_tool—_invokesimplified to a plainawait fn(**args); removediscoroutinefunctioncheck,asyncio.to_threadbranch, and_maybe_await.async def; sync lambdas inmock_toolsreplaced with async functions.Summary by CodeRabbit
Release Notes
Breaking Changes
Tests