Trace MCP calls

package com.example.tnsai.docs;

import com.fasterxml.jackson.databind.JsonNode;
import com.tnsai.mcp.client.McpClient;
import com.tnsai.mcp.client.McpTransport;
import com.tnsai.mcp.observability.McpMetricsCollector;
import com.tnsai.mcp.observability.McpSpan;
import com.tnsai.mcp.observability.McpTracer;
import com.tnsai.mcp.observability.McpTracingInterceptor;
import com.tnsai.mcp.observability.TracingMcpClient;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;

/** Compilation contract for mcp/observability.md. */
public final class McpObservabilityExample {

    private McpObservabilityExample() {
    }

    static Map<String, Object> traceCalls(McpTransport transport) throws IOException {
        McpMetricsCollector metrics = new McpMetricsCollector();
        TracingMcpClient client = new TracingMcpClient(
            McpClient.create(transport), McpTracer.logging(), metrics);   // spans log at DEBUG

        client.connect();
        client.initialize();
        try {
            client.listTools();
            JsonNode content = client.executeTool("search", Map.of("q", "docs"));
            System.out.println(content);
        } finally {
            client.disconnect();
        }
        return metrics.getSummary();       // invocation counts, average ms, error counts, active sessions
    }

    static McpClient traceMessages(McpTransport transport) {
        // Every client, bridge or agent built on this transport gets one span per message.
        return McpClient.create(new McpTracingInterceptor(transport, McpTracer.logging()));
    }

    /** A tracer that hands each finished span to your own sink. */
    public static final class RecordingTracer implements McpTracer {

        public record FinishedSpan(String traceId, String spanId, String operation,
                                   McpSpan.SpanStatus status, Map<String, Object> attributes,
                                   long durationNanos) {
        }

        private final AtomicLong nextSpanId = new AtomicLong();
        private final Consumer<FinishedSpan> sink;

        public RecordingTracer(Consumer<FinishedSpan> sink) {
            this.sink = sink;
        }

        @Override
        public McpSpan startSpan(String operationName) {
            return startSpan(operationName, null);
        }

        @Override
        public McpSpan startSpan(String operationName, McpSpan parent) {
            String traceId = parent != null ? parent.getTraceId() : UUID.randomUUID().toString();
            String spanId = Long.toHexString(nextSpanId.incrementAndGet());
            long start = System.nanoTime();
            Map<String, Object> attributes = new HashMap<>();

            return new McpSpan() {
                private SpanStatus status;     // stays null when the wrapper does not set one

                @Override
                public void setAttribute(String key, String value) {
                    attributes.put(key, value);
                }

                @Override
                public void setAttribute(String key, long value) {
                    attributes.put(key, value);
                }

                @Override
                public void setAttribute(String key, double value) {
                    attributes.put(key, value);
                }

                @Override
                public void setStatus(SpanStatus newStatus) {
                    status = newStatus;
                }

                @Override
                public void recordException(Throwable t) {
                    attributes.put("exception", t.getClass().getName() + ": " + t.getMessage());
                }

                @Override
                public void end() {
                    sink.accept(new FinishedSpan(traceId, spanId, operationName, status,
                        Map.copyOf(attributes), System.nanoTime() - start));
                }

                @Override
                public String getTraceId() {
                    return traceId;
                }

                @Override
                public String getSpanId() {
                    return spanId;
                }
            };
        }
    }

    static TracingMcpClient recorded(McpClient client) {
        return new TracingMcpClient(client, new RecordingTracer(System.out::println));
    }
}