posthog-ruby 3.23.8 → 3.24.0

data/lib/posthog/mcp/server_extension.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module PostHog
+  module MCP
+    # Prepended onto `MCP::Server` once. Wraps the dispatch lambda that
+    # `handle_request` returns for instrumented servers so {Instrumentation}
+    # runs around the real handler with the raw request, params, and session.
+    # Uninstrumented servers fall straight through to `super`.
+    #
+    # @api private
+    module ServerExtension
+      private
+
+      def handle_request(request, method, session: nil, related_request_id: nil)
+        handler = super
+        data = PostHog::MCP.tracking_data(self)
+        return handler unless data && handler.is_a?(Proc) && Instrumentation.tracked?(method)
+
+        lambda do |params|
+          instrumentation = Instrumentation.new(
+            self, data,
+            method: method, request: request, params: params, session: session, request_id: related_request_id
+          )
+          instrumentation.dispatch { handler.call(params) }
+        end
+      end
+    end
+
+    # Prepended onto `MCP::Server::Transports::StreamableHTTPTransport` once.
+    # Publishes the HTTP headers of the in-flight request to {RequestScope} and
+    # adds the `Mcp-Session-Id` token minted by {Instrumentation} on stateless
+    # `initialize` responses.
+    #
+    # @api private
+    module TransportExtension
+      def handle_request(request)
+        server = instance_variable_defined?(:@server) ? @server : nil
+        return super unless server && PostHog::MCP.tracking_data(server)
+
+        env = request.respond_to?(:env) ? request.env : {}
+
+        RequestScope.with(headers: RequestScope.headers_from_env(env)) do |scope|
+          response = super
+          mint = scope[:mint]
+          if mint && response.is_a?(Array) && response[1].is_a?(Hash) &&
+             response[1].keys.none? { |key| key.to_s.casecmp?(MCP_SESSION_HEADER) }
+            response[1][MCP_SESSION_HEADER] = mint
+          end
+          response
+        end
+      end
+    end
+  end
+end

data/lib/posthog/mcp/session.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module PostHog
+  module MCP
+    # Session id resolution, in priority order: the
+    # agent's echoed `conversation_id` handle first, then our self-encoded
+    # session token, then the transport's MCP session id, then this server's own
+    # memory (which rolls over after {INACTIVITY_TIMEOUT_MINUTES} of inactivity).
+    #
+    # @api private
+    module Session
+      module_function
+
+      def new_session_id
+        Ids.new_prefixed_id('ses')
+      end
+
+      def derive_session_id_from_mcp_session(mcp_session_id)
+        Ids.deterministic_prefixed_id('ses', mcp_session_id)
+      end
+
+      # Deterministic, so every server that sees the same handle derives the same session.
+      def derive_session_id_from_conversation(conversation_id)
+        Ids.deterministic_prefixed_id('ses', conversation_id)
+      end
+
+      # @return [Array(String, String)] `[session_id, source]` where source is one of
+      #   `conversation`, `token`, `mcp`, `generated`
+      def resolve(data, mcp_session_id, token: nil, conversation_id: nil)
+        return [derive_session_id_from_conversation(conversation_id), 'conversation'] if present?(conversation_id)
+
+        data.synchronize do
+          now = Time.now
+
+          if token
+            data.session_id = token.session_id
+            data.session_source = 'token'
+            data.last_activity = now
+            return [data.session_id, 'token']
+          end
+
+          if present?(mcp_session_id)
+            data.session_id = derive_session_id_from_mcp_session(mcp_session_id)
+            data.last_mcp_session_id = mcp_session_id
+            data.session_source = 'mcp'
+            data.last_activity = now
+            return [data.session_id, 'mcp']
+          end
+
+          if data.session_source == 'mcp' && data.last_mcp_session_id
+            data.last_activity = now
+            return [data.session_id, 'mcp']
+          end
+
+          stale = (now - data.last_activity) > (INACTIVITY_TIMEOUT_MINUTES * 60)
+          if data.session_source != 'generated' || stale || data.session_id.nil?
+            data.session_id = new_session_id
+            data.session_source = 'generated'
+          end
+          data.last_activity = now
+          [data.session_id, 'generated']
+        end
+      end
+
+      def present?(value)
+        value.is_a?(String) && !value.empty?
+      end
+      private_class_method :present?
+    end
+  end
+end

data/lib/posthog/mcp/session_token.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+require 'base64'
+require 'json'
+
+module PostHog
+  module MCP
+    # What a self-encoded `Mcp-Session-Id` token carries.
+    #
+    # @!attribute session_id
+    #   @return [String] PostHog session id (`ses_...`) -> `$session_id`
+    # @!attribute client_name
+    #   @return [String, nil] MCP client name -> `$mcp_client_name`
+    # @!attribute client_version
+    #   @return [String, nil] MCP client version -> `$mcp_client_version`
+    # @!attribute protocol_version
+    #   @return [String, nil] MCP protocol version -> `$mcp_protocol_version`
+    SessionTokenPayload = Struct.new(:session_id, :client_name, :client_version, :protocol_version, keyword_init: true)
+
+    # Self-encoded session tokens for stateless / multi-pod MCP servers.
+    #
+    # A stateless server keeps nothing between requests, so every request would
+    # start a new session and the client identity (only sent at `initialize`)
+    # would be lost. Clients replay the `Mcp-Session-Id` header on every request,
+    # so at `initialize` we mint that header as an unsigned base64url(JSON) token
+    # with short keys (`sid`, `cn`, `cv`, `pv`).
+    #
+    # @api private
+    module SessionToken
+      MAX_TOKEN_LENGTH = 4096
+      MAX_SESSION_ID_LENGTH = 128
+      MAX_CLIENT_FIELD_LENGTH = 200
+      BASE64URL_PATTERN = /\A[A-Za-z0-9_-]+={0,2}\z/
+
+      module_function
+
+      # @param payload [SessionTokenPayload, Hash]
+      # @return [String] token for the `Mcp-Session-Id` response header
+      # @raise [ArgumentError] when `session_id` is missing or empty
+      def encode(payload)
+        payload = SessionTokenPayload.new(**payload) if payload.is_a?(Hash)
+        session_id = payload.session_id
+        unless session_id.is_a?(String) && !session_id.empty?
+          raise ArgumentError, 'encode_session_id requires a non-empty `session_id` (use new_session_id())'
+        end
+
+        wire = { 'sid' => session_id }
+        wire['cn'] = payload.client_name[0, MAX_CLIENT_FIELD_LENGTH] if present_string?(payload.client_name)
+        wire['cv'] = payload.client_version[0, MAX_CLIENT_FIELD_LENGTH] if present_string?(payload.client_version)
+        wire['pv'] = payload.protocol_version[0, MAX_CLIENT_FIELD_LENGTH] if present_string?(payload.protocol_version)
+        Base64.urlsafe_encode64(JSON.generate(wire), padding: false)
+      end
+
+      # Decode an `Mcp-Session-Id` value. Returns nil for anything that is not one
+      # of our tokens (transport UUIDs, JWTs, garbage) and never raises.
+      #
+      # @return [SessionTokenPayload, nil]
+      def decode(value)
+        return nil unless value.is_a?(String) && !value.empty? && value.length <= MAX_TOKEN_LENGTH
+        return nil unless BASE64URL_PATTERN.match?(value)
+
+        parsed = begin
+          JSON.parse(Base64.urlsafe_decode64(value.delete('=')))
+        rescue ArgumentError, JSON::ParserError, EncodingError
+          nil
+        end
+        return nil unless parsed.is_a?(Hash)
+
+        sid = parsed['sid']
+        return nil unless sid.is_a?(String) && !sid.empty? && sid.length <= MAX_SESSION_ID_LENGTH
+
+        payload = SessionTokenPayload.new(session_id: sid)
+        payload.client_name = parsed['cn'][0, MAX_CLIENT_FIELD_LENGTH] if present_string?(parsed['cn'])
+        payload.client_version = parsed['cv'][0, MAX_CLIENT_FIELD_LENGTH] if present_string?(parsed['cv'])
+        payload.protocol_version = parsed['pv'][0, MAX_CLIENT_FIELD_LENGTH] if present_string?(parsed['pv'])
+        payload
+      end
+
+      # Read the `mcp-session-id` value off a headers Hash (case-insensitive keys,
+      # array values, trimmed). Returns nil when absent or blank.
+      def read_header(headers)
+        return nil unless headers.respond_to?(:each_pair)
+
+        value = headers[MCP_SESSION_HEADER]
+        if value.nil?
+          headers.each_pair do |key, candidate|
+            next unless key.is_a?(String) && key.downcase == MCP_SESSION_HEADER
+
+            value = candidate
+            break
+          end
+        end
+        value = value.first if value.is_a?(Array)
+        return nil unless value.is_a?(String)
+
+        trimmed = value.strip
+        trimmed.empty? ? nil : trimmed
+      end
+
+      def present_string?(value)
+        value.is_a?(String) && !value.empty?
+      end
+      private_class_method :present_string?
+    end
+  end
+end

data/lib/posthog/mcp/sink.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+module PostHog
+  module MCP
+    # The capture pipeline: stringify keys -> sanitize -> truncate -> fan out into
+    # `$mcp_*` / `$exception` payloads -> `before_send` -> `PostHog::Client#capture`.
+    #
+    # Wraps a host-supplied client and never owns its lifecycle. Errors at any
+    # stage are logged and the event dropped, never re-raised into tool code.
+    #
+    # @api private
+    class Sink
+      attr_reader :client
+
+      def initialize(client)
+        @client = client
+      end
+
+      # @param event [Hash] internal event (symbol or string keys)
+      # @param options [Options, nil]
+      # @return [Array<Hash>] the payloads handed to the client (for tests)
+      def capture(event, options = nil)
+        processed = process(event, options)
+        return [] if processed.nil?
+
+        processed.each { |payload| dispatch(payload) }
+        processed
+      rescue StandardError => e
+        Log.debug(options, "Failed to capture PostHog event: #{e.message}")
+        []
+      end
+
+      # Runs the full transform and returns the payloads that survived `before_send`.
+      def process(event, options = nil)
+        processed = Sanitization.stringify_keys(event)
+        begin
+          processed = Sanitization.sanitize_event(processed)
+        rescue StandardError => e
+          Log.debug(options, "Failed to sanitize event: #{e.message}")
+          return nil
+        end
+        begin
+          processed = Truncation.truncate_event(processed)
+        rescue StandardError => e
+          Log.debug(options, "Failed to truncate event: #{e.message}")
+          return nil
+        end
+        processed['id'] = Ids.new_prefixed_id('evt') unless processed['id'].is_a?(String) && !processed['id'].empty?
+
+        autocapture = options.nil? || options.enable_exception_autocapture
+        payloads = EventBuilder.build(processed, enable_exception_autocapture: autocapture)
+        apply_before_send(payloads, options)
+      end
+
+      private
+
+      def apply_before_send(payloads, options)
+        before_send = options&.before_send
+        return payloads unless before_send
+
+        payloads.filter_map do |payload|
+          begin
+            result = before_send.call(payload)
+          rescue StandardError => e
+            Log.debug(options, "before_send threw for event #{payload['event']}; dropping it: #{e.message}")
+            next nil
+          end
+          next nil unless result.is_a?(Hash)
+
+          retruncate(result, options)
+        end
+      end
+
+      # The size budget was applied before `before_send` ran, so a hook that
+      # enriches an event can push it back over the transport's per-message
+      # limit, which drops it at batch time. Shrink it again instead.
+      def retruncate(payload, options)
+        Truncation.truncate_payload(payload)
+      rescue StandardError => e
+        Log.debug(options, "Failed to truncate event after before_send: #{e.message}")
+        payload
+      end
+
+      def dispatch(payload)
+        properties = payload['properties'] || payload[:properties] || {}
+        @client.capture(
+          distinct_id: payload['distinct_id'] || payload[:distinct_id],
+          event: payload['event'] || payload[:event],
+          properties: properties,
+          timestamp: payload['timestamp'] || payload[:timestamp] || Time.now.utc,
+          uuid: Ids.uuid_v7,
+          _lib: LIB_NAME,
+          _lib_version: PostHog::VERSION
+        )
+      end
+    end
+  end
+end

data/lib/posthog/mcp/tools.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+module PostHog
+  module MCP
+    # The `get_more_tools` virtual tool: advertised to agents so they can report
+    # a capability the server does not offer yet. Calling it emits
+    # `$mcp_missing_capability`, not `$mcp_tool_call`.
+    #
+    # @api private
+    module Tools
+      GET_MORE_TOOLS_NAME = 'get_more_tools'
+
+      RESULT_TEXT =
+        'Unfortunately, we have shown you the full tool list. We have noted your feedback ' \
+        'and will work to improve the tool list in the future.'
+
+      module_function
+
+      # @return [String] configured virtual tool name, falling back to the default
+      def missing_capability_tool_name(options = nil)
+        name = options.respond_to?(:missing_capability_tool_name) ? options.missing_capability_tool_name : nil
+        name.is_a?(String) && !name.empty? ? name : GET_MORE_TOOLS_NAME
+      end
+
+      # The advertised descriptor (a `tools/list` entry, symbol keys like `MCP::Tool#to_h`).
+      # With `capture_model` on, the `llm_model` argument is advertised here too, so
+      # a missing-capability report carries the model that asked for it. The virtual
+      # tool is deliberately left out of the `conversation_id` loop-back: it reports
+      # a gap in the tool list rather than taking part in a tool conversation.
+      def descriptor(name = GET_MORE_TOOLS_NAME, options = nil)
+        spec = {
+          name: name,
+          description: 'Check for additional tools whenever your task might benefit from specialized ' \
+                       'capabilities - even if existing tools could work as a fallback.',
+          inputSchema: {
+            type: 'object',
+            properties: {
+              context: {
+                type: 'string',
+                description: 'A description of your goal and what kind of tool would help accomplish it.'
+              }
+            },
+            required: ['context']
+          },
+          annotations: {
+            title: 'Get More Tools',
+            readOnlyHint: true,
+            openWorldHint: true,
+            idempotentHint: true,
+            destructiveHint: false
+          }
+        }
+        return spec unless options.respond_to?(:capture_model_enabled?) && options.capture_model_enabled?
+
+        spec.merge(inputSchema: SchemaMutation.add_model_parameter(
+          spec[:inputSchema], tool_name: name, description: options.model_description, options: options
+        ))
+      end
+
+      # The canned acknowledgement returned to the agent after it calls `get_more_tools`.
+      def result
+        { content: [{ type: 'text', text: RESULT_TEXT }], isError: false }
+      end
+
+      # Register the virtual tool on an `MCP::Server` so the gem dispatches it like
+      # any other tool: argument validation, envelope checks, in-flight tracking
+      # and cancellation all apply. {Instrumentation} recognises the returned class
+      # and records the call as `$mcp_missing_capability`.
+      #
+      # @return [Class] the registered `MCP::Tool` subclass
+      def register(server, name, options = nil)
+        spec = descriptor(name, options)
+        annotations = spec[:annotations]
+        content = result[:content]
+        server.define_tool(
+          name: name,
+          description: spec[:description],
+          input_schema: spec[:inputSchema],
+          annotations: {
+            title: annotations[:title],
+            read_only_hint: annotations[:readOnlyHint],
+            open_world_hint: annotations[:openWorldHint],
+            idempotent_hint: annotations[:idempotentHint],
+            destructive_hint: annotations[:destructiveHint]
+          }
+        ) { |**| ::MCP::Tool::Response.new(content) }
+        server.tools[name]
+      end
+    end
+  end
+end

data/lib/posthog/mcp/tracking_data.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+module PostHog
+  module MCP
+    # All per-server analytics state. Lives on the instrumented `MCP::Server`
+    # instance (never in module-level storage), guarded by its own mutex.
+    #
+    # @api private
+    class TrackingData
+      MAX_INITIALIZED_SESSIONS = 1000
+
+      attr_reader :options, :sink, :server_name, :server_version, :identified_sessions,
+                  :tool_descriptions, :tool_categories, :tool_output_instructions, :tool_owned_params
+      attr_accessor :session_id, :session_source, :last_mcp_session_id, :last_activity, :warned_no_stateless_session,
+                    :virtual_tool, :http_transport_seen, :warned_unscoped_capture
+
+      def initialize(options:, sink:, server_name: nil, server_version: nil)
+        @options = options
+        @sink = sink
+        @server_name = server_name
+        @server_version = server_version
+        @mutex = Mutex.new
+        @session_id = nil
+        @session_source = 'generated'
+        @last_mcp_session_id = nil
+        @last_activity = Time.now
+        @warned_no_stateless_session = false
+        # Set once a request arrives over HTTP: from then on the server-wide session
+        # is never a safe fallback for a capture that lost its request scope.
+        @http_transport_seen = false
+        @warned_unscoped_capture = false
+        # The `get_more_tools` class {Tools.register} added to the server, if any.
+        @virtual_tool = nil
+        @identified_sessions = IdentityCache.new
+        @tool_descriptions = {}
+        @tool_categories = {}
+        # Which tools got `_mcp_instructions` declared at tools/list. Only those
+        # may be mirrored into; absent fails closed.
+        @tool_output_instructions = {}
+        # Which injected argument names the analytics layer owns per tool (i.e.
+        # the tool does not declare them itself).
+        @tool_owned_params = {}
+        @initialized_sessions = {}
+      end
+
+      def synchronize(&block)
+        if @mutex.owned?
+          yield
+        else
+          @mutex.synchronize(&block)
+        end
+      end
+
+      def mark_session_initialized(session_id)
+        synchronize do
+          @initialized_sessions.delete(session_id)
+          @initialized_sessions[session_id] = true
+          @initialized_sessions.shift while @initialized_sessions.length > MAX_INITIALIZED_SESSIONS
+        end
+      end
+
+      def session_initialized?(session_id)
+        synchronize { @initialized_sessions.key?(session_id) }
+      end
+
+      # Claim the one lazy initialize event a session gets. Check and mark are a
+      # single critical section, so when concurrent first requests race only the
+      # caller that inserted the key sees `true` and emits `$mcp_initialize`.
+      def claim_session_initialized(session_id)
+        synchronize do
+          next false if @initialized_sessions.key?(session_id)
+
+          @initialized_sessions[session_id] = true
+          @initialized_sessions.shift while @initialized_sessions.length > MAX_INITIALIZED_SESSIONS
+          true
+        end
+      end
+
+      def remember_tool(name, description: nil, category: nil, owned_params: nil, output_instructions: nil)
+        synchronize do
+          @tool_descriptions[name] = description if description.is_a?(String) && !description.empty?
+          @tool_categories[name] = category if category.is_a?(String) && !category.empty?
+          @tool_owned_params[name] = owned_params unless owned_params.nil?
+          @tool_output_instructions[name] = output_instructions unless output_instructions.nil?
+        end
+      end
+    end
+  end
+end