Skip to content

refactor: extract MCP client initialization logic into LifecyleInitializer #370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

tzolov
Copy link
Contributor

@tzolov tzolov commented Jul 3, 2025

  • Create new LifecyleInitializer class to handle protocol initialization phase
  • Move initialization logic from McpAsyncClient to dedicated initializer
  • Add javadocs for MCP initialization process
  • Implement protocol version negotiation and capability exchange
  • Add exception handling for transport session recovery
  • Include test suite for LifecyleInitializer
  • Simplify McpAsyncClient by delegating initialization responsibilities

Refactor MCP client initialization logic by extracting it from McpAsyncClient into a dedicated LifecyleInitializer class.
This improves separation of concerns, maintainability, and testability of the initialization process.

Untitled-2025-07-05-1728

Motivation and Context

The McpAsyncClient class was becoming too complex with initialization logic mixed in with client operations.
This refactoring addresses several issues:

  • Separation of Concerns: Initialization logic was tightly coupled with client operations, making the code harder to maintain and test
  • Better Error Handling: Centralized exception handling for transport session recovery and re-initialization
  • Improved Testability: Initialization logic can now be tested independently with comprehensive test coverage
  • Cleaner API: The McpAsyncClient interface is now cleaner and focused on its primary responsibilities

How Has This Been Tested?

  • Comprehensive Unit Tests: Added 403 lines of tests in LifecyleInitializerTests.java covering:
    • Successful initialization scenarios
    • Concurrent initialization requests
    • Protocol version negotiation
    • Error handling and recovery
    • Timeout scenarios
    • Transport session exception handling
    • Graceful shutdown procedures
  • Existing Test Suite: All existing McpAsyncClient tests continue to pass, ensuring backward compatibility
  • Integration Testing: Verified that the refactored code works with existing transport implementations (Stdio, SSE)

Breaking Changes

No breaking changes - This is an internal refactoring that maintains full API compatibility. All public methods of McpAsyncClient remain unchanged and behave identically.

Types of changes

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • [] Breaking change (fix or feature that would cause existing functionality to change)
  • Documentation update

Note: While marked as potentially breaking, this is actually a non-breaking internal refactoring.

Checklist

  • I have read the MCP Documentation
  • My code follows the repository's style guidelines
  • New and existing tests pass locally
  • I have added appropriate error handling
  • I have added or updated documentation as needed

Additional context

Implementation Notes:

  • The LifecyleInitializer follows the MCP specification for initialization phase requirements
  • Proper handling of McpTransportSessionNotFoundException with automatic re-initialization
  • Comprehensive JavaDoc documentation explaining the MCP initialization protocol
  • Thread-safe implementation supporting concurrent initialization requests

Files Changed:

  • New: LifecyleInitializer.java (348 lines) - Core initialization logic
  • New: LifecyleInitializerTests.java (403 lines) - Comprehensive test suite
  • Modified: McpAsyncClient.java - Simplified by delegating to LifecyleInitializer

@tzolov tzolov added this to the 0.11.0 milestone Jul 3, 2025
…lizer

- Create new LifecyleInitializer class to handle protocol initialization phase
- Move initialization logic from McpAsyncClient to dedicated initializer
- Add javadocs for MCP initialization process
- Implement protocol version negotiation and capability exchange
- Add exception handling for transport session recovery
- Include  test suite for LifecyleInitializer
- Simplify McpAsyncClient by delegating initialization responsibilities

This refactoring improves separation of concerns and makes the initialization
process more maintainable and testable.

Signed-off-by: Christian Tzolov <[email protected]>
@tzolov tzolov force-pushed the extract-client-initialization-logic branch from 85545f4 to ef140fc Compare July 6, 2025 19:01
* the initialized notification</li>
* </ul>
*/
public class LifecyleInitializer {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
public class LifecyleInitializer {
class LifecyleInitializer {

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's not make it public

* clients and servers.
*/
private final Function<ContextView, McpClientSession> sessionSupplier;
private LifecyleInitializer initializer;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
private LifecyleInitializer initializer;
private final LifecyleInitializer initializer;

* @param operation The operation to execute when the client is initialized
* @return A Mono that completes with the result of the operation
*/
public <T> Mono<T> withIntitialization(String actionName, Function<Initialization, Mono<T>> operation) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
public <T> Mono<T> withIntitialization(String actionName, Function<Initialization, Mono<T>> operation) {
public <T> Mono<T> ensureInitialized(String actionName, Function<Initialization, Mono<T>> operation) {

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think ensureInitialized reflects the functionality It performs. The withIntitialization implies some additional functionality is applied. The ensure implies a validation check.

}

@Test
void shouldTimeoutOnSlowInitialization() {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How long does this test take? Can we use VirtualTimeScheduler for it?

return mockSession;
});

// Start multiple concurrent initializations
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did you intend to make them parallel? You could do that with subscribeOn and using a parallel Scheduler

// re-initialization
// We need to wait a bit for the async re-initialization to start
try {
Thread.sleep(10); // Small delay to allow async processing
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Every time we use sleep we almost certainly introduce flakiness into the test suite, please let's avoid it.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants