Skip to content

Conversation

@ultramancode
Copy link

Implement ToolsRepository for Dynamic Tool Management & Cursor-based Pagination

Closes #746

Summary

This PR introduces the ToolsRepository interface and its default implementation (InMemoryToolsRepository) to enable dynamic, context-aware tool management in McpAsyncServer.

Motivation

Currently, McpAsyncServer only supports static tool registration during initialization. This limits flexibility for scenarios like:

  • Dynamic plugin loading: Tools loaded from external sources at runtime
  • Multi-tenant filtering: Different clients see different tool sets based on their context
  • Permission-based access control: Admin-only tools hidden from regular users
  • Enterprise Scalability: Efficiently handling large-scale toolsets without overwhelming the client or the network. Spec-aligned cursor-based pagination is implemented to minimize payload overhead

Changes

New Files

File Description
ToolsRepository.java Interface for tool management with context-aware operations
InMemoryToolsRepository.java Thread-safe default implementation using ConcurrentHashMap
ToolsListResult.java Record encapsulating list response with pagination cursor

Modified Files

File Changes
McpAsyncServer.java Delegates tool operations to ToolsRepository
McpServerFeatures.java Added toolsRepository field to builder
McpServer.java Builder now supports .toolsRepository() method

Technical Considerations

  • Cursor-based Pagination: The implementation uses cursor tokens (following the MCP spec). The cursor is opaque to clients; while the examples below use a simple numeric offset for demonstration, production implementations can choose stable tokens (e.g., signed/encoded/timestamp-based) to ensure data stability.

  • ToolsListResult Encapsulation: Response is encapsulated in ToolsListResult to cleanly separate the list of tools from the metadata (cursor), matching the MCP tools/list structure.

  • Separation of Concerns: listTools handles visibility ("What tools exist?"), while resolveToolForCall handles executability ("Can I run this?").

Important

Security Note: Hiding a tool in listTools is NOT sufficient for security. Implementations should also verify permissions in resolveToolForCall.

  • Backward Compatibility: Existing .tools() API continues to work unchanged. If no custom repository is provided, InMemoryToolsRepository is automatically used.
  • Thread Safety: InMemoryToolsRepository uses ConcurrentHashMap for safe concurrent access.
  • Extensibility: Developers can implement custom repositories (e.g., Database-backed, Redis-backed) by implementing the ToolsRepository interface.

API Usage

Existing Code (No Changes Required)

McpServer.async(transport)
    .tool(calculatorTool)
    .tool(weatherTool)
    .build();

Custom Repository

McpServer.async(transport)
    .toolsRepository(new MyCustomRepository())
    .build();

Example: Filtering by Client Context and Pagination

Example code
public class RoleBasedToolsRepository implements ToolsRepository {
    private static final String ROLE_ADMIN = "ADMIN";
    private static final String ROLE_USER = "USER";
    private static final String ADMIN_TOOL_PREFIX = "admin-";

    private final List<McpServerFeatures.AsyncToolSpecification> tools = new CopyOnWriteArrayList<>();

    @Override
    public Mono<ToolsListResult> listTools(McpAsyncServerExchange exchange, String cursor) {
        // 1. Extract role from exchange (e.g., from attributes or transport header)
        String role = extractRole(exchange);
        
        // 2. Filter tools based on role
        List<McpServerFeatures.AsyncToolSpecification> visibleTools = tools.stream()
            // Admin tools are visible only to ADMINs
            .filter(spec -> !spec.tool().name().startsWith(ADMIN_TOOL_PREFIX) || ROLE_ADMIN.equals(role))
            .toList();

        // 3. Pagination (With safety guards)
        int pageSize = 10;
        int start = parseCursor(cursor);

        if (start >= visibleTools.size()) {
            return Mono.just(new ToolsListResult(List.of(), null));
        }

        int end = Math.min(start + pageSize, visibleTools.size());

        List<McpSchema.Tool> pageContent = visibleTools.subList(start, end).stream()
            .map(McpServerFeatures.AsyncToolSpecification::tool)
            .toList();

        String nextCursor = (end < visibleTools.size()) ? String.valueOf(end) : null;
            
        return Mono.just(new ToolsListResult(pageContent, nextCursor));
    }

    private int parseCursor(String cursor) {
        if (cursor == null) return 0;
        try {
            return Integer.parseInt(cursor);
        } catch (NumberFormatException e) {
            return 0;
        }
    }

    @Override
    public Mono<McpServerFeatures.AsyncToolSpecification> resolveToolForCall(String name, McpAsyncServerExchange exchange) {
        String role = extractRole(exchange);

        // Security Check: Prevent unauthorized access
        if (name.startsWith(ADMIN_TOOL_PREFIX) && !ROLE_ADMIN.equals(role)) {
            return Mono.empty(); 
        }
        
        return Mono.justOrEmpty(
            tools.stream().filter(t -> t.tool().name().equals(name)).findFirst()
        );
    }

    private String extractRole(McpAsyncServerExchange exchange) {
        // Example: Extract role from custom attributes or transport headers
        // In the verification scenario, this matches the 'X-Role' header
        Object role = exchange.transportContext().get("X-Role");
        return role != null ? role.toString() : ROLE_USER;
    }
}

Pagination Best Practices

  • Stable sorting: Ensure consistent ordering between page requests
  • Context isolation: Cursors should be valid only for the same exchange context (tenant/role)
  • Use null for end: Return null (not empty string) for nextCursor when no more pages exist

Verification Results

The implementation was verified using a sample Spring Boot application to confirm that filtering and pagination work as intended in different role-based scenarios.

Scenario 1: Standard User Access (Filtering & Pagination)

  • Condition: Request with X-Role: USER header.
  • Setup: 5 tools total (2 admin tools, 3 public tools).
  • Behavior:
    • Filtering: Admin tools are correctly hidden from the list.
    • Pagination: Successfully returns visible tools with a valid nextCursor.
  • Result: page1 returned 2 public tools. nextCursor is "2", pointing to the next available tool.
  • View User Request Screenshot (Postman)
mcp_user

Scenario 2: Administrator Access (Full Visibility)

  • Condition: Request with X-Role: ADMIN header.
  • Behavior: All registered tools, including administrative ones, are visible.
  • Result: All 5 tools were returned across pages as expected.
  • View Admin Request Screenshot (Postman)
mcp_admin

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

Checklist

  • I have read the MCP Documentation
  • My code follows the repository's style guidelines (spring-javaformat:apply)
  • New and existing tests pass locally
  • I have added appropriate error handling
  • I have added or updated documentation as needed (Javadoc added)

Additional context

  • Design Decision: The ToolsRepository was introduced to decouple tool storage from the core engine, enabling pluggable and dynamic tool management.
  • Opaque Cursors: While the examples use numeric offsets, the implementation treats cursors as opaque strings to allow for diverse backend implementations (e.g., Redis, DB, or Encrypted tokens).

Important

Security Note: Hiding a tool in listTools is NOT sufficient for security. Implementation developers SHOULD also verify permissions in resolveToolForCall to prevent unauthorized execution.

…or pagination

- Introduce ToolsRepository for runtime tool management
- Support context-aware tool filtering via client exchange
- Implement cursor-based pagination for large toolsets
- Maintain backward compatibility for existing static registration
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.

Support ToolsRepository for Dynamic Tool Management & Cursor-based Pagination Status

1 participant