MCP Context Propagation

package example.mcp.client;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.ai.mcp.client.autoconfigure.NamedClientMcpTransport;
import org.springframework.aop.aspectj.annotation.AspectJProxyFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.stereotype.Component;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonUnwrapped;

import io.micrometer.tracing.CurrentTraceContext;
import io.micrometer.tracing.propagation.Propagator;
import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage;
import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest;

@Component
public final class McpClientTransportsProcessor implements BeanPostProcessor {

    private final McpClientTransportAspect clientTransportAspect;

    public McpClientTransportsProcessor(McpClientTransportAspect clientTransportAspect) {
        this.clientTransportAspect = clientTransportAspect;
    }

    // Multiple transports are part of one bean, so automatic AspectJ weaving doesn't work,
    // manually initialize the proxy on the list items.
    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        if (!(bean instanceof List<?>)) {
            return bean;
        }

        List<?> list = (List<?>) bean;
        if (list.isEmpty()) {
            return bean;
        }
        if (!(list.get(0) instanceof NamedClientMcpTransport)) {
            return bean;
        }

        return list.stream()
                .map(item -> {
                    NamedClientMcpTransport namedTransport = (NamedClientMcpTransport) item;
                    AspectJProxyFactory proxyFactory = new AspectJProxyFactory(namedTransport.transport());
                    proxyFactory.addAspect(clientTransportAspect);
                    return new NamedClientMcpTransport(namedTransport.name(), proxyFactory.getProxy());
                })
                .collect(Collectors.toList());
    }

    @Aspect
    @Component
    public static class McpClientTransportAspect {

        private final CurrentTraceContext currentTraceContext;
        private final Propagator propagator;

        @Autowired
        public McpClientTransportAspect(CurrentTraceContext currentTraceContext, Propagator propagator) {
            this.currentTraceContext = currentTraceContext;
            this.propagator = propagator;
        }

        @Around("execution(* sendMessage(..)) && args(message,..)")
        public Object sendMessage(ProceedingJoinPoint pjp, JSONRPCMessage message) throws Throwable {
            if (!(message instanceof JSONRPCRequest)) {
                return pjp.proceed();
            }

            Map<String, Object> meta = new HashMap<>();
            propagator.inject(currentTraceContext.context(), meta, Map::put);

            if (meta.isEmpty()) {
                return pjp.proceed();
            }

            JSONRPCRequest request = (JSONRPCRequest) message;
            request = new JSONRPCRequest(
                    request.jsonrpc(),
                    request.method(),
                    request.id(),
                    new ParamsWithMeta(request.params(), meta)
            );

            Object[] args = pjp.getArgs();
            args[0] = request;

            return pjp.proceed(args);
        }
    }

    public record ParamsWithMeta(
            @JsonUnwrapped Object params,
            @JsonProperty("_meta") Map<String, Object> meta
    ) {}
}

package example.mcp.server;

import java.util.Map;

import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.aop.aspectj.annotation.AspectJProxyFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import io.micrometer.tracing.CurrentTraceContext;
import io.micrometer.tracing.propagation.Propagator;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest;
import io.modelcontextprotocol.spec.McpServerSession;
import io.opentelemetry.api.GlobalOpenTelemetry;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.propagation.TextMapGetter;
import io.opentelemetry.instrumentation.reactor.v3_1.ContextPropagationOperator;
import reactor.core.publisher.Mono;

@Component
@Aspect
public final class McpServerTransportProviderAspect {

    private final McpServerSessionAspect serverSessionAspect;

    @Autowired
    public McpServerTransportProviderAspect(McpServerSessionAspect serverSessionAspect) {
        this.serverSessionAspect = serverSessionAspect;
    }

    @Around("execution(* io.modelcontextprotocol.spec.McpServerTransportProvider+.setSessionFactory(..)) && args(delegate,..)")
    public Object handle(ProceedingJoinPoint pjp, McpServerSession.Factory delegate) throws Throwable {
        McpServerSession.Factory factory = transport -> {
                AspectJProxyFactory proxyFactory = new AspectJProxyFactory(delegate.create(transport));
            proxyFactory.setProxyTargetClass(true);
            proxyFactory.addAspect(serverSessionAspect);
            return proxyFactory.getProxy();
        };
        Object[] args = pjp.getArgs();
        args[0] = factory;
        return pjp.proceed(args);
    }

    @Aspect
    @Component
    public static class McpServerSessionAspect {
        private final CurrentTraceContext currentTraceContext;
        private final Propagator propagator;

        @Autowired
        public McpServerSessionAspect(CurrentTraceContext currentTraceContext, Propagator propagator) {
            this.currentTraceContext = currentTraceContext;
            this.propagator = propagator;
        }

        @Around("execution(* handle(..)) && args(message,..)")
        @SuppressWarnings({"unchecked"})
        public Object handle(ProceedingJoinPoint pjp, McpSchema.JSONRPCMessage message) throws Throwable {
            if (!(message instanceof JSONRPCRequest)) {
                return pjp.proceed();
            }

            JSONRPCRequest request = (JSONRPCRequest) message;
            if (!(request.params() instanceof Map)) {
                return pjp.proceed();
            }
            Map<String, Object> params = (Map<String, Object>) request.params();

            Object meta = params.get("_meta");
            if (!(meta instanceof Map)) {
                return pjp.proceed();
            }

            // Use OpenTelemetry API directly here
            //   - micrometer only supports starting a new span when extracting, not just extracting
            //   - OTel instrumentation can be used to associate the context with the Mono in a single line.
            //     There may be a simple micrometer or reactor way of doing it but anuraaga doesn't know it.
            Context context = GlobalOpenTelemetry.get().getPropagators().getTextMapPropagator()
                                                 .extract(Context.root(),
                                                          (Map<String, Object>) meta,
                                                          MapGetter.INSTANCE);

            Mono<Void> result = (Mono<Void>) pjp.proceed();
            return ContextPropagationOperator.runWithContext(result, context);
        }
    }

    private enum MapGetter implements TextMapGetter<Map<String, Object>> {
        INSTANCE;

        @Override
        public Iterable<String> keys(Map<String, Object> carrier) {
            return carrier.keySet();
        }

        @Override
        public String get(Map<String, Object> carrier, String key) {
            Object val = carrier.get(key);
            if (!(val instanceof String)) {
                return null;
            }
            return (String) val;
        }
    }
}