/**
 * MCP sampling instrumentation utility.
 *
 * Wraps an MCP client's sampling handler to capture nested LLM calls
 * (sampling/createMessage requests) made during tool execution.
 * The captured data can be reported as child generation spans.
 */
/**
 * Wraps an MCP sampling handler to intercept and measure sampling calls.
 *
 * @param handler - The original sampling handler from the MCP client
 * @param onSamplingCall - Callback invoked with metrics for each sampling call
 * @returns A wrapped handler that transparently captures sampling metrics
 *
 * @example
 * ```ts
 * const samplingCalls: SamplingCallData[] = [];
 * const wrapped = wrapSamplingHandler(
 *   originalHandler,
 *   (data) => samplingCalls.push(data),
 * );
 * // Use `wrapped` as the MCP client's sampling handler
 * // After tool execution, `samplingCalls` contains metrics for all nested LLM calls
 * ```
 */
export function wrapSamplingHandler(handler, onSamplingCall) {
    return async (request) => {
        const startTime = Date.now();
        const response = await handler(request);
        const durationMs = Date.now() - startTime;
        const modelId = response.model ?? request.params.modelPreferences?.hints?.[0]?.name;
        const inputTokens = response.usage?.inputTokens ?? response.usage?.promptTokens;
        const outputTokens = response.usage?.outputTokens ?? response.usage?.completionTokens;
        const reasoningTokens = response.usage?.reasoningTokens;
        const cachedInputTokens = response.usage?.cachedInputTokens;
        onSamplingCall({
            ...(modelId ? { model_id: modelId } : undefined),
            ...(inputTokens != null ? { input_tokens: inputTokens } : undefined),
            ...(outputTokens != null ? { output_tokens: outputTokens } : undefined),
            ...(reasoningTokens != null
                ? { reasoning_tokens: reasoningTokens }
                : undefined),
            ...(cachedInputTokens != null
                ? { cached_input_tokens: cachedInputTokens }
                : undefined),
            duration_ms: durationMs,
        });
        return response;
    };
}
/**
 * Creates a collector that accumulates sampling call data during tool execution.
 * Use with `wrapSamplingHandler` to capture all sampling calls for a tool invocation.
 *
 * @example
 * ```ts
 * const collector = createSamplingCollector();
 * const wrappedHandler = wrapSamplingHandler(handler, collector.collect);
 * // ... execute MCP tool ...
 * const calls = collector.getCalls(); // SamplingCallData[]
 * ```
 */
export function createSamplingCollector() {
    const calls = [];
    return {
        collect: (data) => calls.push(data),
        getCalls: () => [...calls],
        reset: () => {
            calls.length = 0;
        },
    };
}
//# sourceMappingURL=instrumentMcpSampling.js.map