refactor(session): migrate to server-based session management and add tests#9077
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the WalkthroughThe Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant MCPSessionManager
participant MCPServer
Client->>MCPSessionManager: get_session(context_id, connection_params, transport_type)
MCPSessionManager->>MCPSessionManager: Compute server_key from connection_params
alt Existing healthy session for server_key
MCPSessionManager->>MCPSessionManager: Validate session connectivity
alt Session healthy
MCPSessionManager-->>Client: Return existing session
else Session unhealthy
MCPSessionManager->>MCPSessionManager: Cleanup and remove session
MCPSessionManager->>MCPSessionManager: Create new session
MCPSessionManager-->>Client: Return new session
end
else No session for server_key
MCPSessionManager->>MCPSessionManager: Create new session
MCPSessionManager-->>Client: Return new session
end
Note over MCPSessionManager: Periodically cleanup idle sessions
MCPSessionManager->>MCPServer: Communicate as needed via session
Possibly related PRs
Suggested labels
Suggested reviewers
✨ 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 comments)
Other keywords and placeholders
Documentation and Community
|
dosubot
Bot
added
the
size:XXL
coderabbitai
Bot
changed the title
github-actions
Bot
added
refactor
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
src/backend/base/langflow/base/mcp/util.py (2)
528-529: Remove unused import.The
uuidimport on line 528 is not used in this method.- import uuid session_id = f"{server_key}_{len(sessions)}"
549-654: Consider refactoring to reduce code duplication.The
_create_stdio_sessionand_create_sse_sessionmethods have nearly identical implementations. Consider extracting the common session creation pattern into a shared method.Example refactoring approach:
async def _create_session_with_client(self, session_id: str, client_context_manager, transport_type: str): """Generic session creation with any MCP client.""" session_future: asyncio.Future[ClientSession] = asyncio.Future() async def session_task(): try: async with client_context_manager as (read, write): session = ClientSession(read, write) async with session: await session.initialize() session_future.set_result(session) import anyio event = anyio.Event() try: await event.wait() except asyncio.CancelledError: logger.info(f"Session {session_id} is shutting down") except Exception as e: if not session_future.done(): session_future.set_exception(e) # Start the background task task = asyncio.create_task(session_task()) self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) # Wait for session to be ready try: session = await asyncio.wait_for(session_future, timeout=10.0) except asyncio.TimeoutError as timeout_err: # Clean up the failed task if not task.done(): task.cancel() import contextlib with contextlib.suppress(asyncio.CancelledError): await task self._background_tasks.discard(task) msg = f"Timeout waiting for {transport_type.upper()} session {session_id} to initialize" logger.error(msg) raise ValueError(msg) from timeout_err return session, tasksrc/backend/tests/unit/components/data/test_mcp_component.py (2)
155-156: Remove or convert print statement to logging.Debug print statements should not be left in test code.
- # Print tool names for debugging - print(f"Available tools: {tool_names}") + # Log tool names for debugging if needed + # logger.debug(f"Available tools: {tool_names}")
217-217: Remove or convert print statement to logging.Debug print statements should not be left in test code.
- print(f"Available DeepWiki tools: {tool_names}") + # logger.debug(f"Available DeepWiki tools: {tool_names}")
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
src/backend/base/langflow/base/mcp/util.py(7 hunks)src/backend/tests/integration/components/mcp/test_mcp_memory_leak.py(1 hunks)src/backend/tests/unit/components/data/test_mcp_component.py(2 hunks)
🧰 Additional context used
📓 Path-based instructions (8)
src/backend/tests/unit/components/**/*.py
Instructions used from:
Sources:
📄 CodeRabbit Inference Engine
- .cursor/rules/backend_development.mdc
src/backend/**/*.py
Instructions used from:
Sources:
📄 CodeRabbit Inference Engine
- .cursor/rules/backend_development.mdc
src/backend/tests/unit/**/*.py
Instructions used from:
Sources:
📄 CodeRabbit Inference Engine
- .cursor/rules/backend_development.mdc
src/backend/**/*component*.py
Instructions used from:
Sources:
📄 CodeRabbit Inference Engine
- .cursor/rules/icons.mdc
src/backend/**/components/**/*.py
Instructions used from:
Sources:
📄 CodeRabbit Inference Engine
- .cursor/rules/icons.mdc
src/backend/tests/**/*.py
Instructions used from:
Sources:
📄 CodeRabbit Inference Engine
- .cursor/rules/testing.mdc
{src/backend/tests/**/*.py,src/frontend/**/*.test.{ts,tsx,js,jsx},src/frontend/**/*.spec.{ts,tsx,js,jsx},tests/**/*.py}
Instructions used from:
Sources:
📄 CodeRabbit Inference Engine
- .cursor/rules/testing.mdc
{src/backend/tests/**/*.py,tests/**/*.py}
Instructions used from:
Sources:
📄 CodeRabbit Inference Engine
- .cursor/rules/testing.mdc
🧠 Learnings (3)
src/backend/tests/unit/components/data/test_mcp_component.py (17)
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/backend_development.mdc:0-0
Timestamp: 2025-06-30T14:39:17.464Z
Learning: Applies to src/backend/tests/unit/components/**/*.py : Create comprehensive unit tests for all new components
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/backend_development.mdc:0-0
Timestamp: 2025-06-30T14:39:17.464Z
Learning: Applies to src/backend/tests/unit/components/**/*.py : Use ComponentTestBaseWithClient or ComponentTestBaseWithoutClient as base classes for component unit tests
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to src/backend/tests/**/*.py : Inherit from the appropriate ComponentTestBase class ('ComponentTestBase', 'ComponentTestBaseWithClient', or 'ComponentTestBaseWithoutClient') and provide the required fixtures: 'component_class', 'default_kwargs', and 'file_names_mapping' when adding a new component test.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/backend_development.mdc:0-0
Timestamp: 2025-06-30T14:39:17.464Z
Learning: Applies to src/backend/tests/unit/components/**/*.py : Mirror the component directory structure in unit tests under src/backend/tests/unit/components/
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,src/frontend/**/*.test.{ts,tsx,js,jsx},src/frontend/**/*.spec.{ts,tsx,js,jsx},tests/**/*.py} : Create comprehensive unit tests for all new components.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,tests/**/*.py} : Test component configuration updates by asserting changes in build config dictionaries.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,src/frontend/**/*.test.{ts,tsx,js,jsx},src/frontend/**/*.spec.{ts,tsx,js,jsx},tests/**/*.py} : Test both sync and async code paths in components.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,src/frontend/**/*.test.{ts,tsx,js,jsx},src/frontend/**/*.spec.{ts,tsx,js,jsx},tests/**/*.py} : Mock external dependencies appropriately in tests.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,src/frontend/**/*.test.{ts,tsx,js,jsx},src/frontend/**/*.spec.{ts,tsx,js,jsx},tests/**/*.py} : Complex test setups should be commented, and mock usage should be documented within the test code.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,src/frontend/**/*.test.{ts,tsx,js,jsx},src/frontend/**/*.spec.{ts,tsx,js,jsx},tests/**/*.py} : Test component initialization and configuration.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/backend_development.mdc:0-0
Timestamp: 2025-06-30T14:39:17.464Z
Learning: Applies to src/backend/tests/unit/components/**/*.py : Use the client fixture from conftest.py for FastAPI API endpoint tests
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,tests/**/*.py} : Test Langflow's REST API endpoints using the async 'client' fixture and assert correct status codes and response structure.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,tests/**/*.py} : Test components that need external APIs with proper pytest markers such as '@pytest.mark.api_key_required' and '@pytest.mark.no_blockbuster'.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to src/backend/tests/**/*.py : Skip client creation in tests by marking them with '@pytest.mark.noclient' when the 'client' fixture is not needed.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,tests/**/*.py} : Test error handling by monkeypatching internal functions to raise exceptions and asserting correct error responses.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to src/backend/tests/**/*.py : Use the 'client' fixture (an async httpx.AsyncClient) for API tests, as defined in 'src/backend/tests/conftest.py'.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,tests/**/*.py} : Test backward compatibility across Langflow versions by mapping component files to supported versions using 'VersionComponentMapping'.
src/backend/tests/integration/components/mcp/test_mcp_memory_leak.py (14)
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/backend_development.mdc:0-0
Timestamp: 2025-06-30T14:39:17.464Z
Learning: Applies to src/backend/tests/unit/components/**/*.py : Create comprehensive unit tests for all new components
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,tests/**/*.py} : Each test should ensure proper resource cleanup, especially in async fixtures using 'try/finally' and cleanup methods.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,tests/**/*.py} : Be aware of ContextVar propagation in async tests and test both direct event loop execution and 'asyncio.to_thread' scenarios.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,src/frontend/**/*.test.{ts,tsx,js,jsx},src/frontend/**/*.spec.{ts,tsx,js,jsx},tests/**/*.py} : Create comprehensive unit tests for all new components.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,src/frontend/**/*.test.{ts,tsx,js,jsx},src/frontend/**/*.spec.{ts,tsx,js,jsx},tests/**/*.py} : Mock external dependencies appropriately in tests.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,src/frontend/**/*.test.{ts,tsx,js,jsx},src/frontend/**/*.spec.{ts,tsx,js,jsx},tests/**/*.py} : Complex test setups should be commented, and mock usage should be documented within the test code.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,tests/**/*.py} : Test error handling by monkeypatching internal functions to raise exceptions and asserting correct error responses.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to src/backend/tests/**/*.py : Inherit from the appropriate ComponentTestBase class ('ComponentTestBase', 'ComponentTestBaseWithClient', or 'ComponentTestBaseWithoutClient') and provide the required fixtures: 'component_class', 'default_kwargs', and 'file_names_mapping' when adding a new component test.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/backend_development.mdc:0-0
Timestamp: 2025-06-30T14:39:17.464Z
Learning: Applies to src/backend/tests/unit/**/*.py : Test component integration within flows using create_flow, build_flow, and get_build_events utilities
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,src/frontend/**/*.test.{ts,tsx,js,jsx},src/frontend/**/*.spec.{ts,tsx,js,jsx},tests/**/*.py} : Test both sync and async code paths in components.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,tests/**/*.py} : Test Langflow's REST API endpoints using the async 'client' fixture and assert correct status codes and response structure.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/backend_development.mdc:0-0
Timestamp: 2025-06-30T14:39:17.464Z
Learning: Context variables may not propagate correctly in asyncio.to_thread; test both patterns
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/testing.mdc:0-0
Timestamp: 2025-06-30T14:41:58.849Z
Learning: Applies to {src/backend/tests/**/*.py,tests/**/*.py} : Test that operations respect timeout constraints and assert elapsed time is within tolerance.
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/backend_development.mdc:0-0
Timestamp: 2025-06-30T14:39:17.464Z
Learning: Applies to src/backend/base/langflow/components/**/*.py : Use asyncio.create_task for background work in async components and ensure proper cleanup on cancellation
src/backend/base/langflow/base/mcp/util.py (1)
Learnt from: CR
PR: langflow-ai/langflow#0
File: .cursor/rules/backend_development.mdc:0-0
Timestamp: 2025-06-30T14:39:17.464Z
Learning: Applies to src/backend/base/langflow/components/**/*.py : Use asyncio.create_task for background work in async components and ensure proper cleanup on cancellation
🧬 Code Graph Analysis (1)
src/backend/tests/integration/components/mcp/test_mcp_memory_leak.py (1)
src/backend/base/langflow/base/mcp/util.py (5)
MCPSessionManager(387-746)get_session(481-547)cleanup_all(689-722)_get_server_key(449-467)_validate_session_connectivity(469-479)
🔇 Additional comments (17)
src/backend/base/langflow/base/mcp/util.py (7)
33-37: Constants are well-chosen for resource management.The session limits and timeouts are appropriate for preventing resource exhaustion while maintaining good performance.
387-410: Well-implemented initialization with proper background task management.The refactored session manager correctly uses asyncio.create_task for background work and ensures proper cleanup on cancellation, following best practices.
412-422: Robust periodic cleanup implementation.Good error handling and graceful shutdown support.
423-448: Efficient idle session cleanup logic.The implementation correctly removes idle sessions and cleans up empty server entries to prevent memory accumulation.
656-688: Robust cleanup implementation with backward compatibility.The method properly handles both old and new session structures and ensures clean task cancellation.
689-723: Comprehensive cleanup implementation.Excellent handling of all cleanup scenarios including periodic tasks, sessions, and background tasks.
724-747: Good backward compatibility approach.The method clearly documents its purpose and takes the safe approach of cleaning all sessions when the specific session cannot be identified.
src/backend/tests/unit/components/data/test_mcp_component.py (3)
1-57: Well-structured component test following best practices.The test class correctly inherits from ComponentTestBaseWithoutClient and provides all required fixtures as per the testing guidelines.
329-407: Comprehensive integration tests with good coverage.The tests properly verify session management, context setting, and manager sharing as expected from the refactored implementation.
409-454: Good error handling test coverage.The tests properly verify error scenarios and ensure graceful handling of edge cases like multiple disconnect calls.
src/backend/tests/integration/components/mcp/test_mcp_memory_leak.py (7)
1-59: Excellent test setup with subprocess tracking.The fixtures properly handle subprocess tracking and cleanup, ensuring accurate memory leak detection.
61-97: Well-designed test for subprocess leak prevention.The test effectively verifies that session reuse prevents subprocess proliferation, which was the core issue being fixed.
99-126: Good subprocess cleanup verification.The test properly verifies that session cleanup removes subprocesses with appropriate tolerance for timing variations.
128-178: Comprehensive health check and recovery test.The test effectively simulates session failure and verifies recovery, with good backward compatibility handling.
180-223: Effective test for server isolation.The test properly verifies that sessions are isolated between different server configurations.
225-274: Good test coverage for server key generation.This test will help identify the hash stability issue mentioned in the main code review - the test may fail across Python restarts due to hash() instability.
276-365: Well-implemented unit tests with appropriate mocking.The tests effectively verify connectivity validation and cleanup functionality in isolation.
| def _get_server_key(self, connection_params, transport_type: str) -> str: | ||
| """Generate a consistent server key based on connection parameters.""" | ||
| if transport_type == "stdio": | ||
| if hasattr(connection_params, "command"): | ||
| # Include command, args, and environment for uniqueness | ||
| command_str = f"{connection_params.command} {' '.join(connection_params.args or [])}" | ||
| env_str = str(sorted((connection_params.env or {}).items())) | ||
| key_input = f"{command_str}|{env_str}" | ||
| return f"stdio_{hash(key_input)}" | ||
| elif transport_type == "sse": | ||
| if isinstance(connection_params, dict) and "url" in connection_params: | ||
| # Include URL and headers for uniqueness | ||
| url = connection_params["url"] | ||
| headers = str(sorted((connection_params.get("headers", {})).items())) | ||
| key_input = f"{url}|{headers}" | ||
| return f"sse_{hash(key_input)}" | ||
|
|
||
| # Fallback to a generic key | ||
| return f"{transport_type}_{hash(str(connection_params))}" |
There was a problem hiding this comment.
Use stable hashing for server key generation.
Python's hash() function is not stable across Python invocations, which could prevent session reuse after service restarts. Consider using a stable hash like hashlib.md5() or hashlib.sha256().
Apply this diff to fix the hash stability issue:
+import hashlib
+
def _get_server_key(self, connection_params, transport_type: str) -> str:
"""Generate a consistent server key based on connection parameters."""
if transport_type == "stdio":
if hasattr(connection_params, "command"):
# Include command, args, and environment for uniqueness
command_str = f"{connection_params.command} {' '.join(connection_params.args or [])}"
env_str = str(sorted((connection_params.env or {}).items()))
key_input = f"{command_str}|{env_str}"
- return f"stdio_{hash(key_input)}"
+ return f"stdio_{hashlib.md5(key_input.encode()).hexdigest()[:16]}"
elif transport_type == "sse":
if isinstance(connection_params, dict) and "url" in connection_params:
# Include URL and headers for uniqueness
url = connection_params["url"]
headers = str(sorted((connection_params.get("headers", {})).items()))
key_input = f"{url}|{headers}"
- return f"sse_{hash(key_input)}"
+ return f"sse_{hashlib.md5(key_input.encode()).hexdigest()[:16]}"
# Fallback to a generic key
- return f"{transport_type}_{hash(str(connection_params))}"
+ return f"{transport_type}_{hashlib.md5(str(connection_params).encode()).hexdigest()[:16]}"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _get_server_key(self, connection_params, transport_type: str) -> str: | |
| """Generate a consistent server key based on connection parameters.""" | |
| if transport_type == "stdio": | |
| if hasattr(connection_params, "command"): | |
| # Include command, args, and environment for uniqueness | |
| command_str = f"{connection_params.command} {' '.join(connection_params.args or [])}" | |
| env_str = str(sorted((connection_params.env or {}).items())) | |
| key_input = f"{command_str}|{env_str}" | |
| return f"stdio_{hash(key_input)}" | |
| elif transport_type == "sse": | |
| if isinstance(connection_params, dict) and "url" in connection_params: | |
| # Include URL and headers for uniqueness | |
| url = connection_params["url"] | |
| headers = str(sorted((connection_params.get("headers", {})).items())) | |
| key_input = f"{url}|{headers}" | |
| return f"sse_{hash(key_input)}" | |
| # Fallback to a generic key | |
| return f"{transport_type}_{hash(str(connection_params))}" | |
| import hashlib | |
| def _get_server_key(self, connection_params, transport_type: str) -> str: | |
| """Generate a consistent server key based on connection parameters.""" | |
| if transport_type == "stdio": | |
| if hasattr(connection_params, "command"): | |
| # Include command, args, and environment for uniqueness | |
| command_str = f"{connection_params.command} {' '.join(connection_params.args or [])}" | |
| env_str = str(sorted((connection_params.env or {}).items())) | |
| key_input = f"{command_str}|{env_str}" | |
| return f"stdio_{hashlib.md5(key_input.encode()).hexdigest()[:16]}" | |
| elif transport_type == "sse": | |
| if isinstance(connection_params, dict) and "url" in connection_params: | |
| # Include URL and headers for uniqueness | |
| url = connection_params["url"] | |
| headers = str(sorted((connection_params.get("headers", {})).items())) | |
| key_input = f"{url}|{headers}" | |
| return f"sse_{hashlib.md5(key_input.encode()).hexdigest()[:16]}" | |
| # Fallback to a generic key | |
| return f"{transport_type}_{hashlib.md5(str(connection_params).encode()).hexdigest()[:16]}" |
🤖 Prompt for AI Agents
In src/backend/base/langflow/base/mcp/util.py lines 449 to 467, the use of
Python's built-in hash() function for generating server keys is unstable across
Python runs. Replace hash() with a stable hashing method such as hashlib.md5()
or hashlib.sha256(). Convert the key input string to bytes and compute the hash
digest, then use the hex digest as part of the server key to ensure consistent
keys across service restarts.
edwinjosechittilappilly
requested review from
erichare,
jordanrfrazier,
lucaseduoli and
mfortman11
|
@coderabbitai any potential suggestions on the PR and merge readiness. |
|
@coderabbitai and @jordanrfrazier can you review the timeout seconds and check if they need to be changed? |
github-actions
Bot
added
refactor
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chainI'll analyze the PR changes and provide suggestions for merge readiness. Let me examine the key files and implementation details. @edwinjosechittilappilly Based on my analysis of the PR, here are my suggestions and merge readiness assessment: ✅ Merge Readiness: READYThis PR demonstrates excellent engineering practices with comprehensive refactoring, thorough testing, and clear documentation. 🎯 Strengths
📋 Specific Suggestions1. Configuration ConstantsThe new constants ( 2. Error Handling EnhancementConsider adding more specific exception types for different failure scenarios in session management to improve debugging. 3. Monitoring & ObservabilityAdd metrics/logging for:
4. DocumentationThe refactoring changes the behavior significantly - consider updating docstrings to reflect the new server-based approach. 🔍 Code Quality Assessment
🚀 RecommendationAPPROVE for merge - This is a well-executed refactoring that addresses the memory leak issue with proper testing and thoughtful design. The changes are backward compatible and include comprehensive test coverage. The implementation successfully transforms a problematic context-based system into a robust server-based session manager with automatic resource management. 🏁 Script executed: Length of output: 79017 @edwinjosechittilappilly After thoroughly examining the implementation, I can provide a comprehensive assessment: ✅ Final Assessment: EXCELLENT - Ready for MergeThis PR demonstrates exceptional engineering quality with a well-architected solution that thoroughly addresses the memory leak issue. 🔍 Key Implementation Highlights1. Robust Session Management# Server-based session grouping prevents duplicate processes
def _get_server_key(self, connection_params, transport_type: str) -> str:
# Creates consistent keys based on actual server identity
# stdio: hash(command + args + env)
# sse: hash(url + headers)2. Comprehensive Error Handling
3. Session Health & Lifecycleasync def _validate_session_connectivity(self, session) -> bool:
# Quick 3-second health check with list_tools()
# Prevents using dead sessions4. Resource Management
5. Backward Compatibilityasync def _cleanup_session(self, context_id: str):
# Compatibility method for old interface
# Safely handles transition from context-based to server-based🧪 Test Coverage ExcellenceMemory Leak Detection
Error Scenarios
🚀 Production Readiness Indicators✅ Memory Safety: Comprehensive subprocess leak prevention 💡 Minor Recommendations (Post-merge)
🎯 Merge Confidence: 100%This PR:
The implementation quality is production-ready and represents a significant improvement over the previous context-based approach. The memory leak fix is thorough and the session reuse optimization will improve performance. Recommendation: Merge immediately 🚢 🧠 Learnings used |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chainI'll examine the timeout settings in the PR to assess their appropriateness for the MCP session management. 🏁 Script executed: Length of output: 2968 🏁 Script executed: Length of output: 13771 🏁 Script executed: Length of output: 4030 @edwinjosechittilappilly Based on my analysis of the timeout configurations in the MCP session manager, here are my findings and recommendations: Current Timeout Values AnalysisSession Management Timeouts
Connection & Operation Timeouts
Recommendations1. Session Idle Timeout (Currently 5 minutes)Suggestion: Consider reducing to 180-240 seconds (3-4 minutes)
2. Tool Execution Timeout (Currently 30 seconds)Suggestion: Make this configurable or consider reducing to 20-25 seconds
3. Session Cleanup Interval (Currently 1 minute)Suggestion: Keep as-is - well balanced
4. Session Initialization Timeout (Currently 10 seconds)Suggestion: Keep as-is - appropriate for server startup
Code Example for Configurable Timeouts# Consider making these configurable through settings
SESSION_IDLE_TIMEOUT = get_settings_service().settings.mcp_session_idle_timeout or 240 # 4 minutes
TOOL_EXECUTION_TIMEOUT = get_settings_service().settings.mcp_tool_timeout or 25 # 25 seconds
SESSION_CLEANUP_INTERVAL = 60 # Keep as constantMemory Leak ContextGiven that this PR specifically addresses memory leaks, the current timeout values are generally appropriate but could be slightly more aggressive for resource cleanup:
The timeouts are well-structured and comprehensive. The main consideration is whether to prioritize session reuse (longer timeouts) or resource conservation (shorter timeouts) based on your specific deployment needs. (\/) 🧠 Learnings used |
| await self._cleanup_session_by_id(server_key, session_id) | ||
| else: | ||
| # Task is done, clean up | ||
| logger.info(f"Session {session_id} for server {server_key} task is done, cleaning up") |
There was a problem hiding this comment.
Should all of these logs be info or debug -- seems like debug to me
| logger.info( | ||
| f"Maximum sessions reached for server {server_key}, removing oldest session {oldest_session_id}" | ||
| ) | ||
| await self._cleanup_session_by_id(server_key, oldest_session_id) |
There was a problem hiding this comment.
This feels like giving the user the opportunity to unintentionally disrupt expected behavior, by automatically removing sessions past some maximum. I would think that the standard behavior should be to fail to create a new session. But perhaps I need some explaining on what a session is.
There was a problem hiding this comment.
Session is basically a connectin to a subporcess. Each Execution will be created as a new sub process.
Earlier the issue was the old sub process was not closed after execution also not reused.
There was a problem hiding this comment.
Clarified this that the session will be restored on new connection with no disruption to user experience, aside from perhaps needed more time to restore connection.
|
|
||
| # The safest approach is to clean up all sessions since we can't reliably | ||
| # map context_id to the specific session in the new structure | ||
| logger.debug(f"Cleaning up sessions for context_id: {context_id}") |
There was a problem hiding this comment.
This feels like some AI-generated assumption. Do you agree with this?
There was a problem hiding this comment.
True in a way since this backward compatibilty issue never occurs since it get rest on langflow update. but code wise _cleanup_session need not be required hence cleaning all the session.
To validate MCP functionality I also added the tests.
There was a problem hiding this comment.
Catch is that the test will only run for node is available in the env.
github-actions
Bot
added
refactor
github-actions
Bot
added
refactor
|
Flaky tests: |
Added comprehensive unit tests for MCP utility functions, session management, header validation, and client classes in test_mcp_util.py. Updated MCPStdioClient and MCPSseClient disconnect methods for clearer session cleanup logic. Refactored test_mcp_component.py to remove redundant and duplicated tests, consolidating coverage in the new test suite.
github-actions
Bot
added
refactor
…flow into mcp_memory_leak
Codecov Report❌ Patch coverage is
❌ Your project status has failed because the head coverage (2.69%) is below the target coverage (10.00%). You can increase the head coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #9077 +/- ##
==========================================
+ Coverage 33.26% 33.51% +0.24%
==========================================
Files 1194 1194
Lines 56032 56104 +72
Branches 5279 5279
==========================================
+ Hits 18641 18804 +163
+ Misses 37331 37240 -91
Partials 60 60
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
github-actions
Bot
added
refactor
|
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
… tests (#9077) * update MCP Tests * [autofix.ci] apply automated fixes * Update util.py * [autofix.ci] apply automated fixes * Refactor MCP session manager for better configurability and cleanup (#9176) * Add log rotation and header validation features Introduces support for log rotation via the LANGFLOW_LOG_ROTATION environment variable and CLI/config options, with documentation updates. Adds header validation and sanitization for MCP connections, ensuring RFC 7230 compliance and security. Frontend and backend now support passing custom headers for MCP servers. Includes extensive new and updated unit tests for header handling, MCP utilities, and log rotation. * Add unit tests for MCP utilities and update disconnect logic Added comprehensive unit tests for MCP utility functions, session management, header validation, and client classes in test_mcp_util.py. Updated MCPStdioClient and MCPSseClient disconnect methods for clearer session cleanup logic. Refactored test_mcp_component.py to remove redundant and duplicated tests, consolidating coverage in the new test suite. * [autofix.ci] apply automated fixes * Update test_mcp_memory_leak.py * Update util.py --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
… tests (#9077) * update MCP Tests * [autofix.ci] apply automated fixes * Update util.py * [autofix.ci] apply automated fixes * Refactor MCP session manager for better configurability and cleanup (#9176) * Add log rotation and header validation features Introduces support for log rotation via the LANGFLOW_LOG_ROTATION environment variable and CLI/config options, with documentation updates. Adds header validation and sanitization for MCP connections, ensuring RFC 7230 compliance and security. Frontend and backend now support passing custom headers for MCP servers. Includes extensive new and updated unit tests for header handling, MCP utilities, and log rotation. * Add unit tests for MCP utilities and update disconnect logic Added comprehensive unit tests for MCP utility functions, session management, header validation, and client classes in test_mcp_util.py. Updated MCPStdioClient and MCPSseClient disconnect methods for clearer session cleanup logic. Refactored test_mcp_component.py to remove redundant and duplicated tests, consolidating coverage in the new test suite. * [autofix.ci] apply automated fixes * Update test_mcp_memory_leak.py * Update util.py --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
… tests (#9077) * update MCP Tests * [autofix.ci] apply automated fixes * Update util.py * [autofix.ci] apply automated fixes * Refactor MCP session manager for better configurability and cleanup (#9176) * Add log rotation and header validation features Introduces support for log rotation via the LANGFLOW_LOG_ROTATION environment variable and CLI/config options, with documentation updates. Adds header validation and sanitization for MCP connections, ensuring RFC 7230 compliance and security. Frontend and backend now support passing custom headers for MCP servers. Includes extensive new and updated unit tests for header handling, MCP utilities, and log rotation. * Add unit tests for MCP utilities and update disconnect logic Added comprehensive unit tests for MCP utility functions, session management, header validation, and client classes in test_mcp_util.py. Updated MCPStdioClient and MCPSseClient disconnect methods for clearer session cleanup logic. Refactored test_mcp_component.py to remove redundant and duplicated tests, consolidating coverage in the new test suite. * [autofix.ci] apply automated fixes * Update test_mcp_memory_leak.py * Update util.py --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
5 participants
This PR fixes the Memory leak, adds the necessary tests for the memory leak.
Also the MCP Test suite has been improved.
Summary by CodeRabbit
Bug Fixes
Tests