Mandate user id for MCP oauth2 authentication#873
Mandate user id for MCP oauth2 authentication#873rapids-bot[bot] merged 17 commits intoNVIDIA:developfrom
Conversation
Signed-off-by: Binfeng Xu <binengx@nvidia.com>
Signed-off-by: Binfeng Xu <binengx@nvidia.com>
Signed-off-by: Binfeng Xu <binengx@nvidia.com>
Signed-off-by: Anuradha Karuppiah <anuradhak@nvidia.com>
Signed-off-by: Anuradha Karuppiah <anuradhak@nvidia.com>
Signed-off-by: Anuradha Karuppiah <anuradhak@nvidia.com>
Signed-off-by: Anuradha Karuppiah <anuradhak@nvidia.com>
Signed-off-by: Anuradha Karuppiah <anuradhak@nvidia.com>
WalkthroughRequires explicit user/session identification for MCP tool and OAuth flows: auth provider now rejects empty Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Tool as MCPToolClient
participant Ctx as Context (cookies/auth cb)
participant Base as MCPBaseClient
participant Auth as MCPOAuth2Provider
Note over Tool,Auth: Tool call with authentication enabled
Tool->>Tool: _get_session_id(ctx)
Tool->>Ctx: read cookies / auth callback
Ctx-->>Tool: session_id or None
Tool->>Base: access auth_provider
Base-->>Tool: Auth provider instance
alt session_id present
Tool->>Auth: authenticate(user_id=session_id)
Auth-->>Tool: tokens/headers
Tool->>Tool: proceed with tool call using headers
else no session_id
Tool-->>Tool: raise RuntimeError("User is not authorized to call the tool")
end
sequenceDiagram
autonumber
participant Caller as Any Caller
participant Auth as MCPOAuth2Provider
Note over Caller,Auth: Authentication guard
Caller->>Auth: authenticate(user_id)
alt user_id is empty/None
Auth-->>Caller: RuntimeError("User is not authorized to call the tool")
else valid user_id
Auth-->>Caller: proceed with OAuth2 flow
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested labels
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ 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. 🧪 Early access (Sonnet 4.5): enabledWe are currently testing the Sonnet 4.5 model, which is expected to improve code review quality. However, this model may lead to increased noise levels in the review comments. Please disable the early access features if the noise level causes any inconvenience. Note:
Comment |
AnuradhaKaruppiah
changed the title
Signed-off-by: Anuradha Karuppiah <anuradhak@nvidia.com>
|
/ok to test 535c221 |
AnuradhaKaruppiah
added
non-breaking
coderabbitai
bot
added
breaking
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py (2)
123-124: Clarify the session-based tool call behavior in the comment.The updated comment states "Tool call requests should use the session id" but doesn't mention what happens if session_id is None. Consider making this more explicit.
Apply this diff to improve the comment clarity:
- # Tool call requests should use the session id + # Tool call requests must use the session id; if None, auth will fail downstream user_id = session_id
634-638: Consider centralizing the authorization error message.The error message "User is not authorized to call the tool" appears in multiple locations:
auth_provider.pyline 326client_base.pyline 636- Tests validating this message
While the duplication is minimal, consider defining a constant in a shared location (e.g.,
exceptions.py) to ensure consistency if the message needs to change in the future.Example approach:
In
packages/nvidia_nat_mcp/src/nat/plugins/mcp/exceptions.py:# Authorization error messages UNAUTHORIZED_TOOL_CALL_MSG = "User is not authorized to call the tool"Then import and use this constant in both files.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/auth_provider.py(1 hunks)packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py(4 hunks)tests/nat/mcp/test_mcp_auth_provider.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (8)
**/*.{py,yaml,yml}
📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)
**/*.{py,yaml,yml}: Configure response_seq as a list of strings; values cycle per call, and [] yields an empty string.
Configure delay_ms to inject per-call artificial latency in milliseconds for nat_test_llm.
Files:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/auth_provider.pypackages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.pytests/nat/mcp/test_mcp_auth_provider.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)
**/*.py: Programmatic use: create TestLLMConfig(response_seq=[...], delay_ms=...), add with builder.add_llm("", cfg).
When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.) and call the framework’s method (e.g., ainvoke, achat, call).
**/*.py: In code comments/identifiers use NAT abbreviations as specified: nat for API namespace/CLI, nvidia-nat for package name, NAT for env var prefixes; do not use these abbreviations in documentation
Follow PEP 20 and PEP 8; run yapf with column_limit=120; use 4-space indentation; end files with a single trailing newline
Run ruff check --fix as linter (not formatter) using pyproject.toml config; fix warnings unless explicitly ignored
Respect naming: snake_case for functions/variables, PascalCase for classes, UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: use bare raise to re-raise; log with logger.error() when re-raising to avoid duplicate stack traces; use logger.exception() when catching without re-raising
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise and ending with a period; surround code entities with backticks
Validate and sanitize all user input, especially in web or CLI interfaces
Prefer httpx with SSL verification enabled by default and follow OWASP Top-10 recommendations
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile or mprof before optimizing; cache expensive computations with functools.lru_cache or external cache; leverage NumPy vectorized operations when beneficial
Files:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/auth_provider.pypackages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.pytests/nat/mcp/test_mcp_auth_provider.py
packages/*/src/**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Importable Python code inside packages must live under packages//src/
Files:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/auth_provider.pypackages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py
{src/**/*.py,packages/*/src/**/*.py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
All public APIs must have Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Files:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/auth_provider.pypackages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.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:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/auth_provider.pypackages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.pytests/nat/mcp/test_mcp_auth_provider.py
packages/**/*
⚙️ CodeRabbit configuration file
packages/**/*: - This directory contains optional plugin packages for the toolkit, each should contain apyproject.tomlfile. - Thepyproject.tomlfile should declare a dependency onnvidia-nator another package with a name starting
withnvidia-nat-. This dependency should be declared using~=<version>, and the version should be a two
digit version (ex:~=1.0).
- Not all packages contain Python code, if they do they should also contain their own set of tests, in a
tests/directory at the same level as thepyproject.tomlfile.
Files:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/auth_provider.pypackages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py
tests/**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Unit tests reside under tests/ and should use markers defined in pyproject.toml (e.g., integration)
Files:
tests/nat/mcp/test_mcp_auth_provider.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/mcp/test_mcp_auth_provider.py
{tests/**/*.py,examples/*/tests/**/*.py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
{tests/**/*.py,examples/*/tests/**/*.py}: Use pytest (with pytest-asyncio for async); name test files test_*.py; test functions start with test_; extract repeated code into fixtures; fixtures must set name in decorator and be named with fixture_ prefix
Mock external services with pytest_httpserver or unittest.mock; do not hit live endpoints
Mark expensive tests with @pytest.mark.slow or @pytest.mark.integration
Files:
tests/nat/mcp/test_mcp_auth_provider.py
🧬 Code graph analysis (2)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py (4)
src/nat/authentication/interfaces.py (2)
AuthProviderBase(29-72)config(45-54)src/nat/builder/context.py (2)
Context(115-299)user_auth_callback(269-286)packages/nvidia_nat_mcp/src/nat/plugins/mcp/exceptions.py (1)
MCPError(30-43)packages/nvidia_nat_mcp/src/nat/plugins/mcp/exception_handler.py (1)
convert_to_mcp_error(114-144)
tests/nat/mcp/test_mcp_auth_provider.py (2)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/auth_provider.py (2)
MCPOAuth2Provider(283-398)authenticate(313-332)src/nat/authentication/oauth2/oauth2_auth_code_flow_provider.py (1)
authenticate(80-127)
🪛 Ruff (0.13.1)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/auth_provider.py
326-326: Avoid specifying long messages outside the exception class
(TRY003)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py
630-630: Do not catch blind exception: Exception
(BLE001)
⏰ 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 (5)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/auth/auth_provider.py (1)
324-327: LGTM! User ID validation correctly enforces authorization requirement.The early guard prevents tool calls without an authorized user, aligning with the PR objective to mandate user_id for MCP OAuth2 authentication.
Note: The static analysis hint (TRY003) suggests moving the error message to a custom exception class. However, since this message must match the error handling in
client_base.py(line 636) and tests, keeping it inline is acceptable for now. Consider creating a dedicated exception class (e.g.,MCPAuthorizationError) in a future refactor to centralize this error message.tests/nat/mcp/test_mcp_auth_provider.py (1)
297-306: LGTM! Test correctly validates user_id requirement.The test comprehensively covers both scenarios where user_id is missing:
- No user_id argument provided
- Explicit
user_id=NoneBoth cases correctly assert the expected RuntimeError message, ensuring the authentication guard works as intended.
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py (3)
184-186: LGTM! Clean property exposure for auth_provider.The property correctly exposes the internal
_auth_providerattribute, enabling downstream code to access the authentication provider when needed. The type hint properly indicates the optional nature of the provider.
592-616: Well-structured session extraction helper.The
_get_session_idmethod centralizes session ID extraction logic and handles multiple concerns:
- Lazily sets the auth callback from context
- Extracts session_id from cookies
- Falls back to default_user_id when allowed
The implementation is clean and follows the single responsibility principle.
628-638: Authorization check correctly enforces user requirement for authenticated tool calls.The logic properly raises an authorization error when authentication is enabled but no session_id is available, preventing unauthorized tool access. The error message matches the one used in
auth_provider.pyline 326, maintaining consistency.Regarding the static analysis hint (BLE001) on line 630: The bare
Exceptioncatch is acceptable here because:
- The exception is immediately discarded and session_id is set to None
- The subsequent authorization check (lines 635-638) handles the missing session_id case
- Any context access failure shouldn't block the authorization decision
AnuradhaKaruppiah
removed
breaking
Signed-off-by: Anuradha Karuppiah <anuradhak@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py (1)
592-616: Good consolidation of session extraction logic.The helper method centralizes session_id extraction from context and correctly handles the auth callback propagation. The fallback to
default_user_idis properly gated by theallow_default_user_id_for_tool_callsconfiguration flag.Consider enhancing the docstring to document the return value semantics more explicitly:
def _get_session_id(self) -> str | None: """ Get the session id from the context. + + Returns: + Session ID from context cookies, or default_user_id if allowed by config, or None. """
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py(4 hunks)
🧰 Additional context used
📓 Path-based instructions (6)
**/*.{py,yaml,yml}
📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)
**/*.{py,yaml,yml}: Configure response_seq as a list of strings; values cycle per call, and [] yields an empty string.
Configure delay_ms to inject per-call artificial latency in milliseconds for nat_test_llm.
Files:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)
**/*.py: Programmatic use: create TestLLMConfig(response_seq=[...], delay_ms=...), add with builder.add_llm("", cfg).
When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.) and call the framework’s method (e.g., ainvoke, achat, call).
**/*.py: In code comments/identifiers use NAT abbreviations as specified: nat for API namespace/CLI, nvidia-nat for package name, NAT for env var prefixes; do not use these abbreviations in documentation
Follow PEP 20 and PEP 8; run yapf with column_limit=120; use 4-space indentation; end files with a single trailing newline
Run ruff check --fix as linter (not formatter) using pyproject.toml config; fix warnings unless explicitly ignored
Respect naming: snake_case for functions/variables, PascalCase for classes, UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: use bare raise to re-raise; log with logger.error() when re-raising to avoid duplicate stack traces; use logger.exception() when catching without re-raising
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise and ending with a period; surround code entities with backticks
Validate and sanitize all user input, especially in web or CLI interfaces
Prefer httpx with SSL verification enabled by default and follow OWASP Top-10 recommendations
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile or mprof before optimizing; cache expensive computations with functools.lru_cache or external cache; leverage NumPy vectorized operations when beneficial
Files:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py
packages/*/src/**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Importable Python code inside packages must live under packages//src/
Files:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py
{src/**/*.py,packages/*/src/**/*.py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
All public APIs must have Python 3.11+ type hints on parameters and return values; prefer typing/collections.abc abstractions; use typing.Annotated when useful
Files:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.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:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py
packages/**/*
⚙️ CodeRabbit configuration file
packages/**/*: - This directory contains optional plugin packages for the toolkit, each should contain apyproject.tomlfile. - Thepyproject.tomlfile should declare a dependency onnvidia-nator another package with a name starting
withnvidia-nat-. This dependency should be declared using~=<version>, and the version should be a two
digit version (ex:~=1.0).
- Not all packages contain Python code, if they do they should also contain their own set of tests, in a
tests/directory at the same level as thepyproject.tomlfile.
Files:
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py
🧬 Code graph analysis (1)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py (4)
src/nat/authentication/interfaces.py (2)
AuthProviderBase(29-72)config(45-54)src/nat/builder/context.py (2)
Context(115-299)user_auth_callback(269-286)packages/nvidia_nat_mcp/src/nat/plugins/mcp/exceptions.py (1)
MCPError(30-43)packages/nvidia_nat_mcp/src/nat/plugins/mcp/exception_handler.py (1)
convert_to_mcp_error(114-144)
🪛 Ruff (0.13.1)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py
631-631: Do not catch blind exception: Exception
(BLE001)
🔇 Additional comments (3)
packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py (3)
184-186: LGTM: Clean property exposure for auth provider.The read-only property correctly exposes the authentication provider to child components (MCPToolClient) for authorization checks.
635-639: Excellent enforcement of mandatory user authorization.The authorization check correctly implements the PR objective: when authentication is enabled, tool calls now require a valid session_id. The error message is clear, and the conversion to
MCPErrorensures consistent error handling across the MCP client.
122-124: Clear distinction between tool call and non-tool call authentication.The comment accurately describes the behavior: tool call requests use only the session_id (no fallback), while non-tool calls can fall back to
default_user_id. This distinction supports the PR objective of mandating user_id for tool calls.
|
/ok to test 6aa4b6d |
Signed-off-by: Anuradha Karuppiah <anuradhak@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/nat/mcp/test_mcp_client_base.py (2)
554-728: Consider extracting repeated mock setup into pytest fixtures.The test class contains significant code duplication across test methods. The setup of
mock_session,mock_parent_client,mock_metadata, andMCPToolClientinstance is repeated in each test method.Per coding guidelines, frequently repeated code should be extracted into pytest fixtures. Consider creating fixtures for common setup:
@pytest.fixture(name="mock_tool_client_no_auth") def fixture_mock_tool_client_no_auth(): """Fixture providing MCPToolClient with no auth provider.""" mock_session = MagicMock() mock_parent_client = MagicMock() mock_parent_client.auth_provider = None return MCPToolClient( session=mock_session, parent_client=mock_parent_client, tool_name="test_tool", tool_description="Test tool" ) @pytest.fixture(name="mock_tool_client_with_auth") def fixture_mock_tool_client_with_auth(): """Fixture providing MCPToolClient with auth provider configured.""" mock_session = MagicMock() mock_parent_client = MagicMock() mock_auth_provider = MagicMock() mock_auth_config = MagicMock() mock_auth_provider.config = mock_auth_config mock_parent_client.auth_provider = mock_auth_provider return MCPToolClient( session=mock_session, parent_client=mock_parent_client, tool_name="test_tool", tool_description="Test tool" ), mock_auth_config @pytest.fixture(name="mock_context_with_cookies") def fixture_mock_context_with_cookies(): """Fixture providing mock context metadata with cookies.""" mock_metadata = MagicMock() mock_metadata.cookies = {"nat-session": "test-session-123"} return mock_metadataThen simplify tests like:
def test_get_session_id_from_cookies(self, mock_tool_client_no_auth, mock_context_with_cookies): """Test that session_id is correctly extracted from cookies.""" from nat.builder.context import Context as _Ctx with patch.object(_Ctx, 'get') as mock_ctx_get: mock_ctx_get.return_value.metadata = mock_context_with_cookies session_id = mock_tool_client_no_auth._get_session_id() assert session_id == "test-session-123"As per coding guidelines.
554-728: Consider adding test coverage for user_auth_callback interaction.Based on the relevant code snippets (client_base.py lines 534-564),
_get_session_id()retrieves and setsuser_auth_callbackon the parent client when present:auth_callback = _Ctx.get().user_auth_callback if auth_callback and self._parent_client: self._parent_client.set_user_auth_callback(auth_callback)Consider adding a test case to verify this side effect:
def test_get_session_id_sets_auth_callback_when_present(self): """Test that user_auth_callback is set on parent_client when present.""" from nat.builder.context import Context as _Ctx from nat.plugins.mcp.client_base import MCPToolClient mock_session = MagicMock() mock_parent_client = MagicMock() mock_parent_client.auth_provider = None mock_parent_client.set_user_auth_callback = MagicMock() tool_client = MCPToolClient( session=mock_session, parent_client=mock_parent_client, tool_name="test_tool", tool_description="Test tool" ) mock_metadata = MagicMock() mock_metadata.cookies = {"nat-session": "test-session-123"} mock_auth_callback = MagicMock() with patch.object(_Ctx, 'get') as mock_ctx_get: mock_ctx_instance = mock_ctx_get.return_value mock_ctx_instance.metadata = mock_metadata mock_ctx_instance.user_auth_callback = mock_auth_callback session_id = tool_client._get_session_id() # Verify callback was set mock_parent_client.set_user_auth_callback.assert_called_once_with(mock_auth_callback) assert session_id == "test-session-123"
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
examples/MCP/simple_auth_mcp/configs/config-mcp-auth-jira.yml(1 hunks)tests/nat/mcp/test_mcp_client_base.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (8)
**/*.{py,yaml,yml}
📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)
**/*.{py,yaml,yml}: Configure response_seq as a list of strings; values cycle per call, and [] yields an empty string.
Configure delay_ms to inject per-call artificial latency in milliseconds for nat_test_llm.
Files:
tests/nat/mcp/test_mcp_client_base.pyexamples/MCP/simple_auth_mcp/configs/config-mcp-auth-jira.yml
**/*.py
📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)
**/*.py: Programmatic use: create TestLLMConfig(response_seq=[...], delay_ms=...), add with builder.add_llm("", cfg).
When retrieving the test LLM wrapper, use builder.get_llm(name, wrapper_type=LLMFrameworkEnum.) and call the framework’s method (e.g., ainvoke, achat, call).
**/*.py: In code comments/identifiers use NAT abbreviations as specified: nat for API namespace/CLI, nvidia-nat for package name, NAT for env var prefixes; do not use these abbreviations in documentation
Follow PEP 20 and PEP 8; run yapf with column_limit=120; use 4-space indentation; end files with a single trailing newline
Run ruff check --fix as linter (not formatter) using pyproject.toml config; fix warnings unless explicitly ignored
Respect naming: snake_case for functions/variables, PascalCase for classes, UPPER_CASE for constants
Treat pyright warnings as errors during development
Exception handling: use bare raise to re-raise; log with logger.error() when re-raising to avoid duplicate stack traces; use logger.exception() when catching without re-raising
Provide Google-style docstrings for every public module, class, function, and CLI command; first line concise and ending with a period; surround code entities with backticks
Validate and sanitize all user input, especially in web or CLI interfaces
Prefer httpx with SSL verification enabled by default and follow OWASP Top-10 recommendations
Use async/await for I/O-bound work; profile CPU-heavy paths with cProfile or mprof before optimizing; cache expensive computations with functools.lru_cache or external cache; leverage NumPy vectorized operations when beneficial
Files:
tests/nat/mcp/test_mcp_client_base.py
tests/**/*.py
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Unit tests reside under tests/ and should use markers defined in pyproject.toml (e.g., integration)
Files:
tests/nat/mcp/test_mcp_client_base.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/mcp/test_mcp_client_base.py
{tests/**/*.py,examples/*/tests/**/*.py}
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
{tests/**/*.py,examples/*/tests/**/*.py}: Use pytest (with pytest-asyncio for async); name test files test_*.py; test functions start with test_; extract repeated code into fixtures; fixtures must set name in decorator and be named with fixture_ prefix
Mock external services with pytest_httpserver or unittest.mock; do not hit live endpoints
Mark expensive tests with @pytest.mark.slow or @pytest.mark.integration
Files:
tests/nat/mcp/test_mcp_client_base.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/mcp/test_mcp_client_base.pyexamples/MCP/simple_auth_mcp/configs/config-mcp-auth-jira.yml
**/*.{yaml,yml}
📄 CodeRabbit inference engine (.cursor/rules/nat-test-llm.mdc)
In workflow/config YAML, set llms.._type: nat_test_llm to stub responses.
Files:
examples/MCP/simple_auth_mcp/configs/config-mcp-auth-jira.yml
**/configs/**
📄 CodeRabbit inference engine (.cursor/rules/general.mdc)
Configuration files consumed by code must be stored next to that code in a configs/ folder
Files:
examples/MCP/simple_auth_mcp/configs/config-mcp-auth-jira.yml
examples/**/*
⚙️ CodeRabbit configuration file
examples/**/*: - This directory contains example code and usage scenarios for the toolkit, at a minimum an example should
contain a README.md or file README.ipynb.
- If an example contains Python code, it should be placed in a subdirectory named
src/and should
contain apyproject.tomlfile. Optionally, it might also contain scripts in ascripts/directory.- If an example contains YAML files, they should be placed in a subdirectory named
configs/. - If an example contains sample data files, they should be placed in a subdirectory nameddata/, and should
be checked into git-lfs.
Files:
examples/MCP/simple_auth_mcp/configs/config-mcp-auth-jira.yml
🧬 Code graph analysis (1)
tests/nat/mcp/test_mcp_client_base.py (2)
src/nat/builder/context.py (1)
Context(115-299)packages/nvidia_nat_mcp/src/nat/plugins/mcp/client_base.py (3)
MCPToolClient(535-665)auth_provider(185-186)_get_session_id(592-616)
🔇 Additional comments (2)
examples/MCP/simple_auth_mcp/configs/config-mcp-auth-jira.yml (1)
46-46: LGTM! Backward-compatible default preserves existing behavior.The
allow_default_user_id_for_tool_callsconfiguration with a default oftrueappropriately maintains backward compatibility while enabling the new user isolation feature when explicitly disabled.tests/nat/mcp/test_mcp_client_base.py (1)
557-727: LGTM! Comprehensive test coverage for session_id extraction logic.The test class thoroughly validates
_get_session_idbehavior across various scenarios:
- Cookie-based session extraction
- Fallback to default_user_id with proper authorization checks
- Edge cases with missing cookies, missing nat-session, and no auth provider
The tests follow pytest conventions with clear naming, proper mocking, and the arrange-act-assert pattern.
|
/ok to test cf09ea5 |
|
/merge |
3 participants
Description
NAT Oauth2 authenticate abstraction allows headless tool use without user-id. For MCP authentication user-id is crucial for cache maintanence and is mandated for tool calls.
Sample config:
If default_user_id is disallowed for tool calls curl to the server endpoint will fail.
By Submitting this PR I confirm:
Summary by CodeRabbit
Bug Fixes
Tests
Chores