tsikol 0.1.0

data/lib/tsikol/server.rb

@@ -0,0 +1,660 @@
+module Tsikol
+  class Server
+    include Lifecycle
+    include Health
+    
+    attr_reader :name, :version
+    
+    def initialize(name:, version: "1.0.0")
+      @name = name
+      @version = version
+      @tools = {}
+      @resources = {}
+      @prompts = {}
+      @tool_instances = {}
+      @resource_instances = {}
+      @prompt_instances = {}
+      @server_capabilities = {}
+      @log_level = :info
+      @stdio_transport = nil
+      @sampling_handler = nil
+      @middleware_stack = nil
+      @error_handler = ErrorHandler.new(self)
+      
+      # Initialize health monitoring
+      initialize_health_monitoring
+    end
+    
+    # Store transport reference for sending notifications
+    def set_transport(transport)
+      @stdio_transport = transport
+    end
+    
+    def capabilities(&block)
+      instance_eval(&block) if block_given?
+    end
+    
+    def logging(enabled = true)
+      @server_capabilities[:logging] = {} if enabled
+    end
+    
+    def completion(enabled = true)
+      @server_capabilities[:completion] = {} if enabled
+    end
+    
+    def sampling(enabled = true, &block)
+      @server_capabilities[:sampling] = {} if enabled
+      @sampling_handler = block if block_given?
+    end
+    
+    # Alternative method name for clarity
+    def on_sampling(&block)
+      sampling(true, &block)
+    end
+    
+    # Middleware support
+    def use(middleware_class, *args, **kwargs, &block)
+      @middleware_stack ||= MiddlewareStack.new(method(:handle_message_direct))
+      @middleware_stack.use(middleware_class, *args, **kwargs, &block)
+    end
+    
+    # Register class-based instances
+    def register_tool_instance(tool_instance, name: nil)
+      tool_name = name || tool_instance.class.tool_name
+      @tool_instances[tool_name] = tool_instance
+      
+      # Extract completions if any
+      tool_instance.parameters.each do |param_name, config|
+        if config[:completion]
+          @completions ||= {}
+          @completions["tool:#{tool_name}:#{param_name}"] = config[:completion]
+          @server_capabilities[:completion] ||= {}
+        end
+      end
+    end
+    
+    def register_resource_instance(resource_instance, uri: nil)
+      resource_uri = uri || resource_instance.uri
+      @resource_instances[resource_uri] = resource_instance
+    end
+    
+    def register_prompt_instance(prompt_instance, name: nil)
+      prompt_name = name || prompt_instance.class.prompt_name
+      @prompt_instances[prompt_name] = prompt_instance
+      
+      # Extract completions if any
+      prompt_instance.arguments.each do |arg_name, config|
+        if config[:completion]
+          @completions ||= {}
+          @completions["prompt:#{prompt_name}:#{arg_name}"] = config[:completion]
+          @server_capabilities[:completion] ||= {}
+        end
+      end
+    end
+    
+    def tool(name, description: nil, &block)
+      properties, required = extract_parameters_as_schema(block)
+      
+      @tools[name] = {
+        description: description || "Tool: #{name}",
+        handler: block,
+        properties: properties,
+        required: required
+      }
+    end
+    
+    def resource(uri, description: nil, &block)
+      @resources[uri] = {
+        description: description || "Resource: #{uri}",
+        handler: block
+      }
+    end
+    
+    def prompt(name, description: nil, &block)
+      arguments = extract_parameters_as_arguments(block)
+      
+      @prompts[name] = {
+        description: description || "Prompt: #{name}",
+        handler: block,
+        arguments: arguments
+      }
+    end
+    
+    # Define completions for prompts or resources
+    def completion_for(type, name, argument = nil, &block)
+      @completions ||= {}
+      key = "#{type}:#{name}"
+      key += ":#{argument}" if argument
+      @completions[key] = block
+      
+      # Auto-enable completion capability
+      @server_capabilities[:completion] ||= {}
+    end
+    
+    # Logging DSL method
+    def log(level, message, data: nil, logger: nil)
+      # Auto-enable logging capability when first used
+      @server_capabilities[:logging] ||= {}
+      # Define log level hierarchy
+      levels = [:debug, :info, :notice, :warning, :error, :critical, :alert, :emergency]
+      current_level_index = levels.index(@log_level) || 1
+      message_level_index = levels.index(level) || 1
+      
+      # Only log if message level is >= current log level
+      return if message_level_index < current_level_index
+      
+      # Send notification if transport is available
+      if @stdio_transport
+        notification = {
+          jsonrpc: "2.0",
+          method: "notifications/message",
+          params: {
+            level: level.to_s,
+            message: message
+          }
+        }
+        
+        notification[:params][:data] = data if data
+        notification[:params][:logger] = logger if logger
+        
+        @stdio_transport.send_notification(notification)
+      end
+    end
+    
+    def handle_message(raw_message)
+      message = JSON.parse(raw_message)
+      
+      # Use middleware stack if configured
+      if @middleware_stack
+        @middleware_stack.call(message)
+      else
+        handle_message_direct(message)
+      end
+    rescue JSON::ParserError => e
+      # For parse errors, we can't send a response without an ID
+      log :error, "Parse error: #{e.message}"
+      nil
+    rescue => e
+      # Try to send error response if we have an ID
+      if message && message["id"]
+        error_response(message["id"], :internal_error, "Internal error: #{e.message}")
+      else
+        log :error, "Internal error without request ID: #{e.message}"
+        nil
+      end
+    end
+    
+    # Direct message handler (without middleware)
+    def handle_message_direct(message)
+      start_time = Time.now
+      @metrics.increment(:requests_total)
+      @metrics.set(:active_requests, @metrics.get(:active_requests) + 1)
+      
+      response = case message["method"]
+      when "initialize"
+        handle_initialize(message)
+      when "tools/list"
+        handle_tools_list(message)
+      when "tools/call"
+        handle_tool_call(message)
+      when "resources/list"
+        handle_resources_list(message)
+      when "resources/read"
+        handle_resource_read(message)
+      when "prompts/list"
+        handle_prompts_list(message)
+      when "prompts/get"
+        handle_prompt_get(message)
+      when "logging/setLevel"
+        handle_logging_set_level(message)
+      when "completion/complete"
+        handle_completion_complete(message)
+      when "sampling/createMessage"
+        handle_sampling_create_message(message)
+      when "ping"
+        handle_ping(message)
+      else
+        error_response(message["id"], :method_not_found, "Method not found: #{message['method']}")
+      end
+      
+      # Track metrics
+      duration = (Time.now - start_time) * 1000 # Convert to ms
+      @metrics.record(:response_time, duration)
+      @metrics.set(:active_requests, @metrics.get(:active_requests) - 1)
+      
+      if response && response[:error]
+        @metrics.increment(:errors_total)
+      end
+      
+      response
+    end
+    
+    private
+    
+      # Extract parameters for tool schemas (properties format)
+      def extract_parameters_as_schema(block)
+        properties = {}
+        required = []
+        
+        block.parameters.each do |type, param_name|
+          if type == :keyreq  # required keyword argument
+            properties[param_name.to_s] = {
+              type: "string",
+              description: "Required parameter: #{param_name}"
+            }
+            required << param_name.to_s
+          elsif type == :key  # optional keyword argument
+            properties[param_name.to_s] = {
+              type: "string", 
+              description: "Optional parameter: #{param_name}"
+            }
+          end
+        end
+        
+        [properties, required]
+      end
+      
+      # Extract parameters for prompt arguments format
+      def extract_parameters_as_arguments(block)
+        block.parameters.map do |type, param_name|
+          if type == :keyreq || type == :key
+            {
+              name: param_name.to_s,
+              description: "#{type == :keyreq ? 'Required' : 'Optional'} parameter: #{param_name}",
+              required: type == :keyreq
+            }
+          end
+        end.compact
+      end
+      
+      def handle_initialize(message)
+        # Build capabilities based on what's available
+        capabilities = {}
+        capabilities[:tools] = {} if @tools.any?
+        capabilities[:resources] = {} if @resources.any?
+        capabilities[:prompts] = {} if @prompts.any?
+        
+        # Add declared server capabilities
+        capabilities.merge!(@server_capabilities)
+        
+        {
+          jsonrpc: "2.0",
+          id: message["id"],
+          result: {
+            protocolVersion: Tsikol::PROTOCOL_VERSION,
+            capabilities: capabilities,
+            serverInfo: {
+              name: @name,
+              version: @version
+            }
+          }
+        }
+      end
+      
+      def handle_tools_list(message)
+        # Combine both DSL tools and class-based tools
+        tools = []
+        
+        # DSL-style tools
+        @tools.each do |name, tool|
+          tools << {
+            name: name,
+            description: tool[:description],
+            inputSchema: {
+              type: "object",
+              properties: tool[:properties],
+              required: tool[:required]
+            }
+          }
+        end
+        
+        # Class-based tools
+        @tool_instances.each do |name, instance|
+          tools << instance.to_mcp
+        end
+      
+        {
+          jsonrpc: "2.0",
+          id: message["id"],
+          result: { tools: tools }
+        }
+    end
+    
+    def handle_tool_call(message)
+      tool_name = message.dig("params", "name")
+      arguments = message.dig("params", "arguments") || {}
+      
+      # Run before tool hook
+      run_before_tool_hook(tool_name, arguments) if respond_to?(:run_before_tool_hook)
+      
+      # Track tool metrics
+      @metrics.increment("tools:#{tool_name}:calls")
+      start_time = Time.now
+      
+      # Try class-based tool first
+      if (tool_instance = @tool_instances[tool_name])
+        begin
+          kwargs = arguments.transform_keys(&:to_sym)
+          result = @error_handler.wrap_tool_execution(tool_name) do
+            tool_instance.execute(**kwargs)
+          end
+        rescue => e
+          @metrics.increment("tools:#{tool_name}:errors")
+          return error_response(message["id"], :internal_error, "Tool error: #{e.message}")
+        end
+      elsif (tool = @tools[tool_name])
+        # Fall back to DSL-style tool
+        begin
+          result = @error_handler.wrap_tool_execution(tool_name) do
+            if tool[:handler].parameters.any?
+              kwargs = arguments.transform_keys(&:to_sym)
+              tool[:handler].call(**kwargs)
+            else
+              tool[:handler].call
+            end
+          end
+        rescue => e
+          @metrics.increment("tools:#{tool_name}:errors")
+          return error_response(message["id"], :internal_error, "Tool error: #{e.message}")
+        end
+      else
+        return error_response(message["id"], :invalid_params, "Unknown tool: #{tool_name}")
+      end
+      
+      # Track execution time
+      duration = (Time.now - start_time) * 1000
+      @metrics.record("tools:#{tool_name}:duration", duration)
+      
+      # Run after tool hook
+      run_after_tool_hook(tool_name, arguments, result) if respond_to?(:run_after_tool_hook)
+      
+      {
+        jsonrpc: "2.0",
+        id: message["id"],
+        result: {
+          content: [
+            {
+              type: "text",
+              text: result.to_s
+            }
+          ]
+        }
+      }
+    end
+    
+    def handle_resources_list(message)
+      resources = []
+      
+      # DSL-style resources
+      @resources.each do |uri, resource|
+        resources << {
+          uri: uri,
+          name: uri,
+          description: resource[:description]
+        }
+      end
+      
+      # Class-based resources
+      @resource_instances.each do |uri, instance|
+        resources << instance.to_mcp
+      end
+      
+      {
+        jsonrpc: "2.0",
+        id: message["id"],
+        result: { resources: resources }
+      }
+    end
+    
+    def handle_resource_read(message)
+      uri = message.dig("params", "uri")
+      
+      # Try class-based resource first
+      if (resource_instance = @resource_instances[uri])
+        content = resource_instance.read
+      elsif (resource = @resources[uri])
+        # Fall back to DSL-style resource
+        content = resource[:handler].call
+      else
+        return error_response(message["id"], :invalid_params, "Unknown resource: #{uri}")
+      end
+      
+      {
+        jsonrpc: "2.0",
+        id: message["id"],
+        result: {
+          contents: [
+            {
+              uri: uri,
+              mimeType: "text/plain",
+              text: content.to_s
+            }
+          ]
+        }
+      }
+    end
+    
+    def handle_prompts_list(message)
+      prompts = []
+      
+      # DSL-style prompts
+      @prompts.each do |name, prompt|
+        prompts << {
+          name: name,
+          description: prompt[:description],
+          arguments: prompt[:arguments]
+        }
+      end
+      
+      # Class-based prompts
+      @prompt_instances.each do |name, instance|
+        prompts << instance.to_mcp
+      end
+      
+      {
+        jsonrpc: "2.0",
+        id: message["id"],
+        result: { prompts: prompts }
+      }
+    end
+    
+    def handle_prompt_get(message)
+      name = message.dig("params", "name")
+      arguments = message.dig("params", "arguments") || {}
+      
+      # Try class-based prompt first
+      if (prompt_instance = @prompt_instances[name])
+        kwargs = arguments.transform_keys(&:to_sym)
+        messages = prompt_instance.get_messages(**kwargs)
+        
+        return {
+          jsonrpc: "2.0",
+          id: message["id"],
+          result: {
+            messages: messages
+          }
+        }
+      elsif (prompt = @prompts[name])
+        # Fall back to DSL-style prompt
+        kwargs = arguments.transform_keys(&:to_sym)
+        text = prompt[:handler].call(**kwargs)
+      else
+        return error_response(message["id"], :invalid_params, "Unknown prompt: #{name}")
+      end
+      
+      {
+        jsonrpc: "2.0", 
+        id: message["id"],
+        result: {
+          messages: [
+            {
+              role: "user",
+              content: {
+                type: "text",
+                text: text
+              }
+            }
+          ]
+        }
+      }
+    end
+    
+    def handle_logging_set_level(message)
+      level = message.dig("params", "level")
+      
+      # Validate level
+      valid_levels = %w[debug info notice warning error critical alert emergency]
+      unless valid_levels.include?(level)
+        return error_response(message["id"], :invalid_params, "Invalid log level: #{level}")
+      end
+      
+      @log_level = level.to_sym
+      
+      {
+        jsonrpc: "2.0",
+        id: message["id"],
+        result: {}
+      }
+    end
+    
+    def handle_completion_complete(message)
+      params = message["params"] || {}
+      ref = params["ref"]
+      argument = params["argument"]
+      
+      unless ref
+        return error_response(message["id"], :invalid_params, "Missing ref parameter")
+      end
+      
+      # Build completion key
+      completion_key = case ref["type"]
+      when "ref/prompt"
+        if argument
+          "prompt:#{ref['name']}:#{argument['name']}"
+        else
+          "prompt:#{ref['name']}"
+        end
+      when "ref/resource"
+        "resource:#{ref['uri']}"
+      else
+        return error_response(message["id"], :invalid_params, "Unknown ref type: #{ref['type']}")
+      end
+      
+      # Get completion handler
+      handler = @completions&.[](completion_key)
+      
+      # Get completions
+      values = if handler
+        begin
+          # Call handler with current value
+          current_value = argument&.dig("value") || ""
+          result = handler.call(current_value)
+          Array(result)
+        rescue => e
+          log :error, "Completion handler error", data: { error: e.message }
+          []
+        end
+      else
+        []
+      end
+      
+      # Limit to 100 results as per spec
+      values = values.first(100)
+      
+      {
+        jsonrpc: "2.0",
+        id: message["id"],
+        result: {
+          completion: {
+            values: values,
+            total: values.size,
+            hasMore: false
+          }
+        }
+      }
+    end
+    
+    def handle_sampling_create_message(message)
+      params = message["params"] || {}
+      
+      # Validate required parameters
+      messages = params["messages"]
+      unless messages && messages.is_a?(Array)
+        return error_response(message["id"], :invalid_params, "Missing or invalid messages parameter")
+      end
+      
+      # Extract parameters
+      model_preferences = params["modelPreferences"] || {}
+      system_prompt = params["systemPrompt"]
+      max_tokens = params["maxTokens"]
+      include_context = params["includeContext"] || "none"
+      temperature = params["temperature"]
+      metadata = params["metadata"] || {}
+      
+      # Build sampling request data
+      sampling_data = {
+        messages: messages,
+        model_preferences: model_preferences,
+        system_prompt: system_prompt,
+        max_tokens: max_tokens,
+        temperature: temperature,
+        metadata: metadata
+      }
+      
+      # Call registered sampling handler if any
+      if @sampling_handler
+        begin
+          result = @sampling_handler.call(sampling_data)
+          
+          # Result should have role and content
+          unless result.is_a?(Hash) && result[:role] && result[:content]
+            raise "Sampling handler must return hash with :role and :content"
+          end
+          
+          {
+            jsonrpc: "2.0",
+            id: message["id"],
+            result: result
+          }
+        rescue => e
+          error_response(message["id"], :internal_error, "Sampling error: #{e.message}")
+        end
+      else
+        # No sampling handler registered
+        error_response(message["id"], :internal_error, "No sampling handler registered. Call 'on_sampling' to register a handler.")
+      end
+    end
+    
+    def handle_ping(message)
+      # Simple ping/pong response
+      {
+        jsonrpc: "2.0",
+        id: message["id"],
+        result: {}
+      }
+    end
+    
+    def error_response(id, code_key, message)
+      # If no ID, we can't send a proper error response
+      return nil if id.nil?
+      
+      code = case code_key
+      when :parse_error then Errors::PARSE_ERROR
+      when :invalid_request then Errors::INVALID_REQUEST
+      when :method_not_found then Errors::METHOD_NOT_FOUND
+      when :invalid_params then Errors::INVALID_PARAMS
+      when :internal_error then Errors::INTERNAL_ERROR
+      else Errors::INTERNAL_ERROR
+      end
+      
+      {
+        jsonrpc: "2.0",
+        id: id,
+        error: {
+          code: code,
+          message: message
+        }
+      }
+    end
+  end
+end