Pull Request Review - PR #489

Summary

This PR addresses two important improvements:

1. Fixed docstring update mechanism in KnowledgeBase search method
2. Added MCP handler cleanup logic to prevent resource leaks

Code Quality Review

✅ Strengths

1. Resource Management: The MCP cleanup logic properly prevents resource leaks by closing task-level handlers
2. Error Handling: Comprehensive error suppression for event loop issues (common in threaded contexts)
3. Selective Cleanup: Smart distinction between agent-level and task-level tools to avoid closing shared resources
4. Canvas Integration: Automatic canvas tool registration improves usability

⚠️ Issues & Concerns

1. KnowledgeBase Docstring Update (Medium Priority)

File: src/upsonic/knowledge_base/knowledge_base.py:174-179

Issue: The docstring update mechanism may not work correctly for bound methods.

if hasattr(self.search, '__func__'):
    self.search.__func__.__doc__ = base_docstring.format(...)

Concern:

• For bound methods, modifying __func__.__doc__ affects the underlying function, but the docstring may not be visible when introspecting the bound method itself
• Consider testing with inspect.getdoc() to verify the docstring is properly accessible
• Alternative approach: Store the formatted docstring separately or use functools.wraps()

Recommendation: Add a unit test that verifies the docstring is properly updated and accessible via inspect.getdoc(kb.search).

2. MCP Handler Cleanup Race Condition (Low-Medium Priority)

File: src/upsonic/agent/pipeline/steps.py:1262-1266

Issue: Set membership check may not work as intended.

agent_tools_set = set(context.agent.tools) if context.agent.tools else set()
for tool in context.task.tools:
    if isinstance(tool, (MCPHandler, MultiMCPHandler)):
        if tool not in agent_tools_set:

Concern:

• Object identity comparison in sets uses __hash__ and __eq__. If MCPHandler doesn't implement these properly, the same logical tool might not match
• If tools are wrapped/proxied during registration, identity checks may fail

Recommendation:

• Verify MCPHandler implements __hash__ and __eq__ based on connection parameters, not object identity
• Consider using IDs or connection strings for comparison instead of object identity
• Add logging in debug mode to show which handlers are being closed

3. Canvas Tools Mutation Side Effect (Low Priority)

File: src/upsonic/agent/agent.py:558-563

Issue: Mutating self.tools has side effects.

if self.canvas:
    canvas_functions = self.canvas.functions()
    for canvas_func in canvas_functions:
        if canvas_func not in final_tools:
            final_tools.append(canvas_func)
    self.tools = final_tools  # ← Modifies original attribute

Concern:

• This mutates the agent's tools attribute after initialization
• May cause confusion if users inspect agent.tools and see tools they didn't explicitly provide
• Could cause issues if _register_agent_tools() is called multiple times

Recommendation:

• Document this behavior clearly in the Agent class docstring
• Consider making this explicit in initialization rather than during tool registration
• Or use a separate self._effective_tools attribute to avoid surprising mutations

4. Removed Comments in vectordb/config.py

File: src/upsonic/vectordb/config.py:117, 145, 195-198

Issue: Removed section divider comments and explanatory comments.

-# ============================================================================
-# Index Tuning Configurations
-# ============================================================================

Concern: These comments improved code organization and readability. Their removal makes the file harder to navigate without clear benefit.

Recommendation: Consider keeping structural comments that help with code navigation, or explain the reasoning for removing them in the PR description.

Security Considerations

✅ No security issues identified. The MCP cleanup logic properly handles exceptions and doesn't expose sensitive data.

Performance Considerations

✅ Positive impact: Proper resource cleanup prevents connection leaks that could accumulate over time with many task executions.

⚠️ Minor concern: The cleanup loop iterates through all task tools for every task completion. For agents with many tools, consider:

• Caching the set of MCP handlers during registration
• Using a dedicated list of task-level MCP handlers instead of filtering at cleanup time

Test Coverage Recommendations

The PR would benefit from:

1. Unit test for docstring update:

   def test_knowledge_base_search_docstring_includes_description():
       kb = KnowledgeBase(name="test", description="Test KB", topics=["topic1"])
       doc = inspect.getdoc(kb.search)
       assert "Test KB" in doc
       assert "topic1" in doc

2. Integration test for MCP cleanup:

   async def test_task_mcp_handlers_cleaned_up_after_execution():
       # Create agent with no tools
       # Create task with MCP handler
       # Execute task
       # Verify handler.close() was called

3. Test for canvas tool integration:

   def test_canvas_tools_automatically_registered():
       canvas = Mock()
       canvas.functions.return_value = [mock_func]
       agent = Direct(canvas=canvas, tools=[other_tool])
       assert mock_func in agent.registered_agent_tools

Best Practices & Suggestions

1. ✅ Good use of defensive programming with try-except blocks
2. ✅ Proper async/await usage in cleanup logic
3. ✅ Debug-mode logging for troubleshooting
4. ⚠️ Consider adding type hints to new code sections for consistency
5. ⚠️ Document the canvas auto-registration behavior in the Agent class docstring

Final Verdict

Recommendation: Approve with minor revisions

This PR fixes real issues (docstring updates and resource leaks), but would benefit from:

1. Verification that the docstring update mechanism works correctly
2. Tests to prevent regression
3. Consideration of the object identity comparison issue in MCP cleanup

The changes are production-ready but could be more robust with the suggested improvements.

Priority Actions:

1. Add unit test for KnowledgeBase docstring update
2. Verify MCPHandler set membership comparison works as intended
3. Document canvas tool auto-registration behavior
4. Consider keeping helpful code organization comments

Great work on addressing resource cleanup! 👍