Client Handler Annotations

/**
 * Marks a method as a handler for logging notifications from MCP servers
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@interface McpLogging {
    /**
     * Client connection name(s) to handle logging from
     * Required - must specify at least one client
     */
    String[] clients();
}

@Component
public class ServerLogHandler {

    @McpLogging(clients = "weather-server")
    public void handleWeatherLogs(LoggingMessageNotification notification) {
        LoggingLevel level = notification.level();
        String logger = notification.logger();
        Object data = notification.data();

        System.out.println("[" + level + "] " + logger + ": " + data);
    }
}

@Component
public class LogProcessor {

    @McpLogging(clients = {"server1", "server2"})
    public void processLogs(LoggingLevel level, String logger, String data) {
        if (level == LoggingLevel.ERROR || level == LoggingLevel.CRITICAL) {
            alertingService.sendAlert("Server error: " + data);
        }

        logStorage.store(logger, level, data);
    }
}

@Component
public class AsyncLogHandler {

    @McpLogging(clients = "production-server")
    public Mono<Void> handleLogsAsync(LoggingMessageNotification notification) {
        return Mono.fromRunnable(() -> {
            // Process logs asynchronously
            logDatabase.saveAsync(notification);
        });
    }
}

/**
 * Marks a method as a handler for LLM sampling requests from MCP servers
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@interface McpSampling {
    /**
     * Client connection name(s) to handle sampling requests from
     * Required - must specify at least one client
     */
    String[] clients();
}

@Component
public class SamplingHandler {

    private final ChatClient chatClient;

    @McpSampling(clients = "agent-server")
    public CreateMessageResult handleSampling(CreateMessageRequest request) {
        List<PromptMessage> messages = request.messages();
        String systemPrompt = request.systemPrompt();

        // Use Spring AI ChatClient to generate response
        String response = chatClient.prompt()
            .system(systemPrompt != null ? systemPrompt : "")
            .user(extractUserMessage(messages))
            .call()
            .content();

        return new CreateMessageResult(
            response,
            Role.ASSISTANT,
            "gpt-4",
            "end_turn"
        );
    }

    private String extractUserMessage(List<PromptMessage> messages) {
        return messages.stream()
            .filter(m -> m.role() == Role.USER)
            .map(m -> ((TextContent) m.content()).text())
            .collect(Collectors.joining("\n"));
    }
}

@Component
public class AdvancedSamplingHandler {

    @McpSampling(clients = {"server1", "server2"})
    public CreateMessageResult handleAdvancedSampling(CreateMessageRequest request) {
        Map<String, Object> modelPrefs = request.modelPreferences();
        Integer maxTokens = request.maxTokens();

        // Configure chat client based on preferences
        var promptBuilder = chatClient.prompt();

        if (request.systemPrompt() != null) {
            promptBuilder.system(request.systemPrompt());
        }

        // Add messages from request
        for (PromptMessage msg : request.messages()) {
            if (msg.role() == Role.USER) {
                promptBuilder.user(extractTextContent(msg));
            }
        }

        // Apply max tokens if specified
        if (maxTokens != null) {
            promptBuilder.options(ChatOptions.builder()
                .maxTokens(maxTokens)
                .build());
        }

        String response = promptBuilder.call().content();

        return new CreateMessageResult(
            response,
            Role.ASSISTANT,
            "gpt-4",
            "end_turn"
        );
    }
}

@Component
public class ReactiveSamplingHandler {

    @McpSampling(clients = "async-server")
    public Mono<CreateMessageResult> handleSamplingAsync(CreateMessageRequest request) {
        return Mono.fromCallable(() -> {
            // Generate response asynchronously
            String response = llmService.generateCompletion(request);

            return new CreateMessageResult(
                response,
                Role.ASSISTANT,
                "gpt-4-turbo",
                "stop_sequence"
            );
        }).subscribeOn(Schedulers.boundedElastic());
    }
}

/**
 * Marks a method as a handler for elicitation requests from MCP servers
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@interface McpElicitation {
    /**
     * Client connection name(s) to handle elicitation requests from
     * Required - must specify at least one client
     */
    String[] clients();
}

@Component
public class ElicitationHandler {

    private final UserInputService userInputService;

    @McpElicitation(clients = "interactive-server")
    public ElicitResult handleElicitation(ElicitRequest request) {
        String message = request.message();
        Map<String, Object> schema = request.requestedSchema();

        // Prompt user for input
        boolean accepted = userInputService.promptUser(message);

        if (accepted) {
            Map<String, Object> data = userInputService.collectInput(schema);
            return new ElicitResult(ElicitResult.Action.ACCEPT, data);
        } else {
            return new ElicitResult(ElicitResult.Action.DECLINE, Map.of());
        }
    }
}

@Component
public class UIElicitationHandler {

    private final ApplicationEventPublisher eventPublisher;

    @McpElicitation(clients = {"form-server", "wizard-server"})
    public ElicitResult handleUIElicitation(ElicitRequest request) {
        // Show UI dialog to user
        ElicitationDialog dialog = new ElicitationDialog(
            request.message(),
            request.requestedSchema()
        );

        ElicitationDialog.Result result = dialog.showAndWait();

        if (result.isAccepted()) {
            return new ElicitResult(
                ElicitResult.Action.ACCEPT,
                result.getData()
            );
        } else if (result.isCancelled()) {
            return new ElicitResult(
                ElicitResult.Action.CANCEL,
                Map.of()
            );
        } else {
            return new ElicitResult(
                ElicitResult.Action.DECLINE,
                Map.of()
            );
        }
    }
}

@Component
public class AsyncElicitationHandler {

    @McpElicitation(clients = "async-interactive-server")
    public Mono<ElicitResult> handleElicitationAsync(ElicitRequest request) {
        return webSocketService.requestUserInput(request.message(), request.requestedSchema())
            .map(userData -> new ElicitResult(ElicitResult.Action.ACCEPT, userData))
            .defaultIfEmpty(new ElicitResult(ElicitResult.Action.DECLINE, Map.of()));
    }
}

/**
 * Marks a method as a handler for progress notifications from MCP servers
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@interface McpProgress {
    /**
     * Client connection name(s) to handle progress notifications from
     * Required - must specify at least one client
     */
    String[] clients();
}

@Component
public class ProgressHandler {

    @McpProgress(clients = "long-running-server")
    public void handleProgress(ProgressNotification notification) {
        String token = notification.progressToken();
        double progress = notification.progress();
        Double total = notification.total();
        String message = notification.message();

        if (total != null) {
            double percentage = (progress / total) * 100;
            System.out.println(token + ": " + percentage + "% - " + message);
        } else {
            System.out.println(token + ": " + message);
        }
    }
}

@Component
public class UIProgressHandler {

    private final Map<String, ProgressBar> progressBars = new ConcurrentHashMap<>();

    @McpProgress(clients = {"worker1", "worker2"})
    public void updateProgressBar(String progressToken, double progress, Double total, String message) {
        ProgressBar bar = progressBars.computeIfAbsent(
            progressToken,
            token -> new ProgressBar()
        );

        if (total != null) {
            bar.setProgress(progress / total);
        }

        if (message != null) {
            bar.setLabel(message);
        }
    }
}

@Component
public class AsyncProgressHandler {

    @McpProgress(clients = "background-processor")
    public Mono<Void> handleProgressAsync(ProgressNotification notification) {
        return Mono.fromRunnable(() -> {
            // Update database or cache with progress info
            progressRepository.updateProgress(
                notification.progressToken(),
                notification.progress(),
                notification.total()
            );
        });
    }
}

/**
 * Marks a method as a handler for tool list change notifications
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@interface McpToolListChanged {
    /**
     * Client connection name(s) to handle tool list changes from
     * Required - must specify at least one client
     */
    String[] clients();
}

@Component
public class ToolChangeHandler {

    private final ToolCacheService toolCache;

    @McpToolListChanged(clients = "dynamic-server")
    public void handleToolListChanged(List<McpSchema.Tool> tools) {
        System.out.println("Tool list updated. Total tools: " + tools.size());

        // Update cache
        toolCache.refresh(tools);

        // Log tool names
        tools.forEach(tool -> System.out.println("  - " + tool.name()));
    }
}

@Component
public class AsyncToolChangeHandler {

    @McpToolListChanged(clients = {"server1", "server2"})
    public Mono<Void> handleToolChanges(List<McpSchema.Tool> tools) {
        return Mono.fromRunnable(() -> {
            // Invalidate caches
            toolCallbackProvider.invalidateCache();

            // Notify UI
            eventBus.publish(new ToolsUpdatedEvent(tools));
        });
    }
}

/**
 * Marks a method as a handler for resource list change notifications
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@interface McpResourceListChanged {
    /**
     * Client connection name(s) to handle resource list changes from
     * Required - must specify at least one client
     */
    String[] clients();
}

@Component
public class ResourceChangeHandler {

    @McpResourceListChanged(clients = "data-server")
    public void handleResourceListChanged(List<McpSchema.Resource> resources) {
        System.out.println("Resource list updated. Total resources: " + resources.size());

        resources.forEach(resource -> {
            System.out.println("  - " + resource.name() + " (" + resource.uri() + ")");
        });

        // Update local resource cache
        resourceCache.update(resources);
    }
}

/**
 * Marks a method as a handler for prompt list change notifications
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@interface McpPromptListChanged {
    /**
     * Client connection name(s) to handle prompt list changes from
     * Required - must specify at least one client
     */
    String[] clients();
}

@Component
public class PromptChangeHandler {

    @McpPromptListChanged(clients = "prompt-server")
    public void handlePromptListChanged(List<McpSchema.Prompt> prompts) {
        System.out.println("Prompt list updated. Total prompts: " + prompts.size());

        prompts.forEach(prompt -> {
            System.out.println("  - " + prompt.name());
            prompt.arguments().forEach(arg ->
                System.out.println("      * " + arg.name() + " (" + arg.required() + ")")
            );
        });

        // Refresh prompt catalog
        promptCatalog.refresh(prompts);
    }
}

@Component
public class ComprehensiveClientHandlers {

    private final Logger log = LoggerFactory.getLogger(getClass());
    private final ChatClient chatClient;
    private final UserInterfaceService uiService;

    // Handle logging from all servers
    @McpLogging(clients = {"server1", "server2", "server3"})
    public void handleAllLogs(LoggingLevel level, String logger, String data) {
        log.info("[{}] {}: {}", level, logger, data);
    }

    // Handle sampling requests
    @McpSampling(clients = "assistant-server")
    public CreateMessageResult handleSampling(CreateMessageRequest request) {
        String response = chatClient.prompt()
            .messages(convertMessages(request.messages()))
            .call()
            .content();

        return new CreateMessageResult(response, Role.ASSISTANT, "gpt-4", "end_turn");
    }

    // Handle user input requests
    @McpElicitation(clients = "interactive-agent")
    public ElicitResult handleUserInput(ElicitRequest request) {
        Map<String, Object> userData = uiService.showDialog(request.message(), request.requestedSchema());
        return new ElicitResult(ElicitResult.Action.ACCEPT, userData);
    }

    // Track progress
    @McpProgress(clients = "batch-processor")
    public void trackProgress(ProgressNotification notification) {
        uiService.updateProgressBar(
            notification.progressToken(),
            notification.progress(),
            notification.total()
        );
    }

    // React to tool changes
    @McpToolListChanged(clients = "dynamic-tools-server")
    public void onToolsChanged(List<McpSchema.Tool> tools) {
        log.info("Tools updated: {} available", tools.size());
        uiService.refreshToolMenu(tools);
    }

    private List<Message> convertMessages(List<PromptMessage> mcpMessages) {
        // Convert MCP messages to Spring AI messages
        return mcpMessages.stream()
            .map(this::convertMessage)
            .collect(Collectors.toList());
    }

    private Message convertMessage(PromptMessage mcpMessage) {
        String content = extractTextContent(mcpMessage.content());
        return switch (mcpMessage.role()) {
            case USER -> new UserMessage(content);
            case ASSISTANT -> new AssistantMessage(content);
            case SYSTEM -> new SystemMessage(content);
        };
    }

    private String extractTextContent(Content content) {
        if (content instanceof TextContent textContent) {
            return textContent.text();
        }
        return "";
    }
}

@Configuration
public class McpClientConfiguration {

    @Bean
    public McpSyncClient assistantServer() {
        return McpClient.sync()
            .transport(SseClientTransport.builder()
                .url("http://assistant-server.example.com/sse")
                .build())
            .build();
    }

    @Bean
    public McpSyncClient dynamicToolsServer() {
        return McpClient.sync()
            .transport(StdioServerTransport.builder()
                .command("dynamic-tools-mcp")
                .build())
            .build();
    }

    // Handler beans are automatically discovered and wired to clients
}

/**
 * Registry of methods annotated with MCP Client annotations
 * Scans and registers handlers for MCP client notifications and requests
 * Location: org.springframework.ai.mcp.annotation.spring.ClientMcpSyncHandlersRegistry
 */
class ClientMcpSyncHandlersRegistry implements SmartInitializingSingleton {
    /**
     * Get MCP client capabilities for a specific client
     * @param clientName The name of the MCP client
     * @return ClientCapabilities declared by annotations for this client
     */
    McpSchema.ClientCapabilities getCapabilities(String clientName);

    /**
     * Invoke the sampling handler for a given MCP client
     * @param name The client name
     * @param samplingRequest The sampling request from the server
     * @return Result of the sampling operation
     */
    McpSchema.CreateMessageResult handleSampling(String name, McpSchema.CreateMessageRequest samplingRequest);

    /**
     * Invoke the elicitation handler for a given MCP client
     * @param name The client name
     * @param elicitationRequest The elicitation request from the server
     * @return Result of the elicitation operation
     */
    McpSchema.ElicitResult handleElicitation(String name, McpSchema.ElicitRequest elicitationRequest);

    /**
     * Invoke all logging handlers for a given MCP client
     * @param name The client name
     * @param loggingMessageNotification The logging notification
     */
    void handleLogging(String name, McpSchema.LoggingMessageNotification loggingMessageNotification);

    /**
     * Invoke all progress handlers for a given MCP client
     * @param name The client name
     * @param progressNotification The progress notification
     */
    void handleProgress(String name, McpSchema.ProgressNotification progressNotification);

    /**
     * Invoke all tool list changed handlers for a given MCP client
     * @param name The client name
     * @param updatedTools The updated list of tools
     */
    void handleToolListChanged(String name, List<McpSchema.Tool> updatedTools);

    /**
     * Invoke all prompt list changed handlers for a given MCP client
     * @param name The client name
     * @param updatedPrompts The updated list of prompts
     */
    void handlePromptListChanged(String name, List<McpSchema.Prompt> updatedPrompts);

    /**
     * Invoke all resource list changed handlers for a given MCP client
     * @param name The client name
     * @param updatedResources The updated list of resources
     */
    void handleResourceListChanged(String name, List<McpSchema.Resource> updatedResources);
}

@Configuration
public class CustomClientHandlerConfig {

    private final ClientMcpSyncHandlersRegistry registry;

    // Check client capabilities
    public void checkCapabilities(String clientName) {
        McpSchema.ClientCapabilities caps = registry.getCapabilities(clientName);

        if (caps.sampling() != null) {
            System.out.println("Client supports sampling");
        }
        if (caps.elicitation() != null) {
            System.out.println("Client supports elicitation");
        }
    }

    // Manually trigger handler (usually done automatically by the framework)
    public void triggerLogging(String clientName, LoggingMessageNotification notification) {
        registry.handleLogging(clientName, notification);
    }
}

/**
 * Registry of methods annotated with MCP Client annotations (async/reactive)
 * Scans and registers reactive handlers for MCP client notifications and requests
 * Location: org.springframework.ai.mcp.annotation.spring.ClientMcpAsyncHandlersRegistry
 */
class ClientMcpAsyncHandlersRegistry implements SmartInitializingSingleton {
    /**
     * Get MCP client capabilities for a specific client
     * @param clientName The name of the MCP client
     * @return ClientCapabilities declared by annotations for this client
     */
    McpSchema.ClientCapabilities getCapabilities(String clientName);

    /**
     * Invoke the sampling handler for a given MCP client (async)
     * @param name The client name
     * @param samplingRequest The sampling request from the server
     * @return Mono wrapping the sampling result
     */
    Mono<McpSchema.CreateMessageResult> handleSampling(String name, McpSchema.CreateMessageRequest samplingRequest);

    /**
     * Invoke the elicitation handler for a given MCP client (async)
     * @param name The client name
     * @param elicitationRequest The elicitation request from the server
     * @return Mono wrapping the elicitation result
     */
    Mono<McpSchema.ElicitResult> handleElicitation(String name, McpSchema.ElicitRequest elicitationRequest);

    /**
     * Invoke all logging handlers for a given MCP client (async)
     * @param name The client name
     * @param loggingMessageNotification The logging notification
     * @return Mono that completes when all handlers finish
     */
    Mono<Void> handleLogging(String name, McpSchema.LoggingMessageNotification loggingMessageNotification);

    /**
     * Invoke all progress handlers for a given MCP client (async)
     * @param name The client name
     * @param progressNotification The progress notification
     * @return Mono that completes when all handlers finish
     */
    Mono<Void> handleProgress(String name, McpSchema.ProgressNotification progressNotification);

    /**
     * Invoke all tool list changed handlers for a given MCP client (async)
     * @param name The client name
     * @param updatedTools The updated list of tools
     * @return Mono that completes when all handlers finish
     */
    Mono<Void> handleToolListChanged(String name, List<McpSchema.Tool> updatedTools);

    /**
     * Invoke all prompt list changed handlers for a given MCP client (async)
     * @param name The client name
     * @param updatedPrompts The updated list of prompts
     * @return Mono that completes when all handlers finish
     */
    Mono<Void> handlePromptListChanged(String name, List<McpSchema.Prompt> updatedPrompts);

    /**
     * Invoke all resource list changed handlers for a given MCP client (async)
     * @param name The client name
     * @param updatedResources The updated list of resources
     * @return Mono that completes when all handlers finish
     */
    Mono<Void> handleResourceListChanged(String name, List<McpSchema.Resource> updatedResources);
}

@Configuration
public class ReactiveClientHandlerConfig {

    private final ClientMcpAsyncHandlersRegistry registry;

    // Check capabilities
    public Mono<Void> processClient(String clientName) {
        McpSchema.ClientCapabilities caps = registry.getCapabilities(clientName);

        // Chain async operations
        return Mono.just(caps)
            .filter(c -> c.sampling() != null)
            .flatMap(c -> testSampling(clientName))
            .then();
    }

    // Manually trigger async handler
    public Mono<Void> triggerProgressUpdate(String clientName, ProgressNotification notification) {
        return registry.handleProgress(clientName, notification);
    }
}