Add observability support when using MCP front end#741
Add observability support when using MCP front end#741rapids-bot[bot] merged 2 commits intoNVIDIA:developfrom
Conversation
- Updated Workflow class to expose exporter_manager through a property for better access. - Modified get_all_exporters method to utilize the new exporter_manager property. - Enhanced register_function_with_mcp to accept a workflow parameter, allowing for observability context during function registration. - Added tests to validate the new functionality and ensure observability context is correctly propagated. Signed-off-by: Matthew Penn <mpenn@nvidia.com>
WalkthroughAdds a public Workflow.exporter_manager property, propagates Workflow into MCP tool registration, and wraps MCP-invoked functions with workflow-driven observability using ContextState and the exporter_manager. Adds tests covering wrapper creation, observability lifecycle, description resolution, and registration. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant U as User/Client
participant MCP as MCP Front End
participant Wrapper as Tool Wrapper
participant WF as Workflow
participant EM as WF.exporter_manager
participant CS as ContextState
participant F as NAT Function
U->>MCP: Invoke tool (name, args)
MCP->>Wrapper: Call registered wrapper(args, workflow)
Wrapper->>WF: access exporter_manager
WF-->>Wrapper: ExporterManager
Wrapper->>EM: start(tool_call_context)
activate EM
EM-->>Wrapper: async enter
Wrapper->>CS: ContextState.get() inside started context
CS-->>Wrapper: Observability context
Wrapper->>F: acall_invoke/ainvoke(args)
F-->>Wrapper: Result / Error
Wrapper->>EM: exit context
deactivate EM
Wrapper-->>MCP: Result / Propagated Error
MCP-->>U: Response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Assessment against linked issues
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. 📜 Recent review detailsConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro 💡 Knowledge Base configuration:
You can enable these sources in your CodeRabbit configuration. 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/nat/front_ends/mcp/tool_converter.py (2)
35-41: Annotate return type of public APIcreate_function_wrapper.Public functions in src/ must have return types. This returns an async callable producing
str.Apply this diff:
-def create_function_wrapper( +def create_function_wrapper( function_name: str, function: FunctionBase, schema: type[BaseModel], is_workflow: bool = False, - workflow: 'Workflow | None' = None, -) : + workflow: 'Workflow | None' = None, +) -> Callable[..., Awaitable[str]]:
75-86: Fix Pydantic v2 required/optional detection.
field.is_requiredin v2 is a callable; using it as a boolean risks misclassifying required fields, breaking tool parameter defaults.Apply this diff:
- for name, field in param_fields.items(): + for name, field in param_fields.items(): # Get the field type and convert to appropriate Python type field_type = field.annotation - - # Add the parameter to our list - parameters.append( - Parameter( - name=name, - kind=Parameter.KEYWORD_ONLY, - default=Parameter.empty if field.is_required else None, - annotation=field_type, - )) + # Determine required/optional robustly across Pydantic versions + is_required_attr = getattr(field, "is_required", None) + required = is_required_attr() if callable(is_required_attr) else bool(is_required_attr) + default = Parameter.empty if required else None + + parameters.append( + Parameter( + name=name, + kind=Parameter.KEYWORD_ONLY, + default=default, + annotation=field_type, + ))
🧹 Nitpick comments (7)
src/nat/builder/workflow.py (1)
89-91: Document the new publicexporter_managerproperty.Add a short Google‑style docstring so downstream users know it returns a copy‑on‑write ExporterManager.
Apply this diff:
@property def exporter_manager(self) -> ExporterManager: - return self._exporter_manager.get() + """Public accessor for the workflow's ExporterManager. + + Returns: + ExporterManager: Copy-on-write manager for the workflow's telemetry exporters. + """ + return self._exporter_manager.get()src/nat/front_ends/mcp/tool_converter.py (3)
118-123: Avoid redundant.get()onworkflow.exporter_manager.
Workflow.exporter_manageralready returns anExporterManagerfrom.get(). Calling.get()again is unnecessary.Apply this diff:
- context_state = ContextState.get() - exporter_manager = workflow.exporter_manager.get() + context_state = ContextState.get() + exporter_manager = workflow.exporter_manager async with exporter_manager.start(context_state=context_state): return await func_call()Note: If you adopt this, update tests that stub
workflow.exporter_manager.get()to assert onworkflow.exporter_manager.start(...)instead.
204-209: Docstring precedence doesn’t match implementation.Code prefers
config.descriptionoverconfig.topic; update the docstring to match.Apply this diff:
- The description is determined using the following precedence: - 1. If the function is a Workflow and has a 'description' attribute, use it. - 2. If the Workflow's config has a 'topic', use it. - 3. If the Workflow's config has a 'description', use it. - 4. If the function is a regular Function, use its 'description' attribute. + The description is determined using the following precedence: + 1. If the function is a Workflow and has a 'description' attribute, use it. + 2. If the Workflow's config has a 'description', use it. + 3. If the Workflow's config has a 'topic', use it. + 4. If the function is a regular Function, use its 'description' attribute.
182-186: Optional: also log to module logger on exceptions.You re‑raise (good). Consider
logger.error("...")alongsidectx.error(...)for centralized logs.tests/nat/front_ends/mcp/test_tool_converter.py (3)
265-283: Mark async tests and assert exporter_manager.start usage.Add marker and strengthen the observability assertion.
Apply this diff:
- @patch('nat.front_ends.mcp.tool_converter.ContextState') - async def test_observability_context_propagation(self, mock_context_state_class): + @pytest.mark.asyncio + @patch('nat.front_ends.mcp.tool_converter.ContextState') + async def test_observability_context_propagation(self, mock_context_state_class): @@ - mock_workflow.exporter_manager.get.assert_called_once() - mock_context_state_class.get.assert_called_once() + mock_workflow.exporter_manager.get.assert_called_once() + mock_context_state_class.get.assert_called_once() + mock_workflow.exporter_manager.get.return_value.start.assert_called_once_with( + context_state=mock_context_state + )
289-306: Mark async test and assert observability start even on error.Strengthen guarantees about lifecycle.
Apply this diff:
- @patch('nat.front_ends.mcp.tool_converter.ContextState') - async def test_error_handling_in_wrapper_execution(self, mock_context_state_class): + @pytest.mark.asyncio + @patch('nat.front_ends.mcp.tool_converter.ContextState') + async def test_error_handling_in_wrapper_execution(self, mock_context_state_class): @@ - mock_workflow.exporter_manager.get.assert_called_once() - mock_context_state_class.get.assert_called_once() + mock_workflow.exporter_manager.get.assert_called_once() + mock_context_state_class.get.assert_called_once() + mock_workflow.exporter_manager.get.return_value.start.assert_called_once_with( + context_state=mock_context_state + )
32-36: Minor: no need to set name in MockChatRequest.
"ChatRequest" in __qualname__already satisfies detection; consider removing to reduce confusion.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (4)
src/nat/builder/workflow.py(2 hunks)src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py(1 hunks)src/nat/front_ends/mcp/tool_converter.py(11 hunks)tests/nat/front_ends/mcp/test_tool_converter.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (8)
tests/**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Unit tests must live under tests/ and use configured markers (e2e, integration, etc.)
Files:
tests/nat/front_ends/mcp/test_tool_converter.py
⚙️ CodeRabbit configuration file
tests/**/*.py: - Ensure that tests are comprehensive, cover edge cases, and validate the functionality of the code. - Test functions should be named using thetest_prefix, using snake_case. - Any frequently repeated code should be extracted into pytest fixtures. - Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture
function being decorated should be named using thefixture_prefix, using snake_case. Example:
@pytest.fixture(name="my_fixture")
def fixture_my_fixture():
pass
Files:
tests/nat/front_ends/mcp/test_tool_converter.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.py: Follow PEP 8/20 style; format with yapf (column_limit=120) and use 4-space indentation; end files with a single newline
Run ruff (ruff check --fix) per pyproject.toml; fix warnings unless explicitly ignored; ruff is linter-only
Use snake_case for functions/variables, PascalCase for classes, and UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: preserve stack traces and avoid duplicate logging
When re-raising exceptions, use bareraiseand log with logger.error(), not logger.exception()
When catching and not re-raising, log with logger.exception() to capture stack trace
Validate and sanitize all user input; prefer httpx with SSL verification and follow OWASP Top‑10
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile/mprof; cache with functools.lru_cache or external cache; leverage NumPy vectorization when beneficial
Files:
tests/nat/front_ends/mcp/test_tool_converter.pysrc/nat/builder/workflow.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/mcp/tool_converter.py
**/tests/**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/tests/**/*.py: Test functions must use the test_ prefix and snake_case
Extract repeated test code into pytest fixtures; fixtures should set name=... in @pytest.fixture and functions named with fixture_ prefix
Mark expensive tests with @pytest.mark.slow or @pytest.mark.integration
Use pytest with pytest-asyncio for async code; mock external services with pytest_httpserver or unittest.mock
Files:
tests/nat/front_ends/mcp/test_tool_converter.py
**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
**/*.{py,sh,md,yml,yaml,toml,ini,json,ipynb,txt,rst}: Every file must start with the standard SPDX Apache-2.0 header; keep copyright years up‑to‑date
All source files must include the SPDX Apache‑2.0 header; do not bypass CI header checks
Files:
tests/nat/front_ends/mcp/test_tool_converter.pysrc/nat/builder/workflow.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/mcp/tool_converter.py
**/*.{py,md}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Never hard‑code version numbers in code or docs; versions are derived by setuptools‑scm
Files:
tests/nat/front_ends/mcp/test_tool_converter.pysrc/nat/builder/workflow.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/mcp/tool_converter.py
**/*
⚙️ CodeRabbit configuration file
**/*: # Code Review Instructions
- Ensure the code follows best practices and coding standards. - For Python code, follow
PEP 20 and
PEP 8 for style guidelines.- Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values.
Example:def my_function(param1: int, param2: str) -> bool: pass- For Python exception handling, ensure proper stack trace preservation:
- When re-raising exceptions: use bare
raisestatements to maintain the original stack trace,
and uselogger.error()(notlogger.exception()) to avoid duplicate stack trace output.- When catching and logging exceptions without re-raising: always use
logger.exception()
to capture the full stack trace information.Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any
words listed in the
ci/vale/styles/config/vocabularies/nat/reject.txtfile, words that might appear to be
spelling mistakes but are listed in theci/vale/styles/config/vocabularies/nat/accept.txtfile are OK.Misc. - All code (except .mdc files that contain Cursor rules) should be licensed under the Apache License 2.0,
and should contain an Apache License 2.0 header comment at the top of each file.
- Confirm that copyright years are up-to date whenever a file is changed.
Files:
tests/nat/front_ends/mcp/test_tool_converter.pysrc/nat/builder/workflow.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/mcp/tool_converter.py
src/**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
src/**/*.py: All importable Python code must live under src/
All public APIs in src/ require Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise with a period; surround code entities with backticks
Files:
src/nat/builder/workflow.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/mcp/tool_converter.py
src/nat/**/*
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Core functionality under src/nat should prioritize backward compatibility when changed
Files:
src/nat/builder/workflow.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/mcp/tool_converter.py
⚙️ CodeRabbit configuration file
This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.
Files:
src/nat/builder/workflow.pysrc/nat/front_ends/mcp/mcp_front_end_plugin_worker.pysrc/nat/front_ends/mcp/tool_converter.py
🧬 Code graph analysis (4)
tests/nat/front_ends/mcp/test_tool_converter.py (3)
src/nat/runtime/session.py (1)
workflow(84-85)src/nat/builder/workflow.py (2)
Workflow(40-155)exporter_manager(90-91)src/nat/front_ends/mcp/tool_converter.py (3)
create_function_wrapper(35-197)get_function_description(200-237)register_function_with_mcp(240-269)
src/nat/builder/workflow.py (2)
tests/nat/observability/test_exporter_manager.py (1)
exporter_manager(121-123)src/nat/observability/exporter_manager.py (1)
ExporterManager(26-335)
src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py (2)
src/nat/front_ends/mcp/tool_converter.py (1)
register_function_with_mcp(240-269)src/nat/runtime/session.py (1)
workflow(84-85)
src/nat/front_ends/mcp/tool_converter.py (3)
src/nat/builder/context.py (1)
ContextState(65-90)src/nat/runtime/session.py (1)
workflow(84-85)src/nat/builder/workflow.py (2)
Workflow(40-155)exporter_manager(90-91)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: CI Pipeline / Check
🔇 Additional comments (4)
src/nat/builder/workflow.py (2)
87-88: Good switch to public accessor in get_all_exporters.Delegating via the new
exporter_manageraccessor is correct and future‑proof.
100-104: LGTM: Runner now receives the workflow’s exporter_manager.This enables tracing for workflow‑scoped runs via MCP without extra plumbing.
src/nat/front_ends/mcp/mcp_front_end_plugin_worker.py (1)
137-140: Correctly propagating workflow into MCP registration.Passing
workflowintoregister_function_with_mcpis the key to enabling observability; looks good.src/nat/front_ends/mcp/tool_converter.py (1)
140-171: No change needed:ainvokeandacall_invokeare both implemented in FunctionBase (src/nat/builder/function.py), withacall_invokedelegating toainvoke.
Signed-off-by: Matthew Penn <mpenn@nvidia.com>
|
/merge |
3 participants
Description
Closes #740
This PR adds observability support for all MCP function and workflow calls. Previously, functions and workflows executed through the MCP front end bypassed the observability layer, missing telemetry and tracing data. This change wraps all MCP executions with the workflow's observability context using
workflow.exporter_manager.get().start(), ensuring consistent monitoring across all NAT execution paths. Includes comprehensive tests covering both success and error scenarios.By Submitting this PR I confirm:
Summary by CodeRabbit
New Features
Improvements
Documentation
Tests