llmemory 0.1.10 → 0.1.11

data/lib/llmemory/mcp/authentication.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module MCP
+    # Rack middleware for token-based authentication
+    # Validates requests against MCP_TOKEN environment variable
+    class Authentication
+      UNAUTHORIZED_RESPONSE = [
+        401,
+        { "Content-Type" => "application/json" },
+        ['{"error":"Unauthorized: Invalid or missing token"}']
+      ].freeze
+
+      def initialize(app, token: nil)
+        @app = app
+        @token = token || ENV["MCP_TOKEN"]
+      end
+
+      def call(env)
+        # If no token is configured, skip authentication
+        return @app.call(env) unless @token
+
+        # Check for valid token
+        return UNAUTHORIZED_RESPONSE unless valid_token?(env)
+
+        @app.call(env)
+      end
+
+      private
+
+      def valid_token?(env)
+        # Check Authorization header (Bearer token)
+        auth_header = env["HTTP_AUTHORIZATION"]
+        if auth_header
+          # Support both "Bearer <token>" and plain "<token>"
+          token = auth_header.sub(/\ABearer\s+/i, "")
+          return true if secure_compare(token, @token)
+        end
+
+        # Check query string parameter
+        query_string = env["QUERY_STRING"] || ""
+        query_params = parse_query_string(query_string)
+        if query_params["token"]
+          return true if secure_compare(query_params["token"], @token)
+        end
+
+        false
+      end
+
+      def parse_query_string(query_string)
+        query_string.split("&").each_with_object({}) do |pair, hash|
+          key, value = pair.split("=", 2)
+          hash[key] = value if key
+        end
+      end
+
+      # Constant-time comparison to prevent timing attacks
+      def secure_compare(a, b)
+        return false if a.nil? || b.nil?
+        return false if a.bytesize != b.bytesize
+
+        l = a.unpack("C*")
+        r = b.unpack("C*")
+        result = 0
+        l.zip(r) { |x, y| result |= x ^ y }
+        result.zero?
+      end
+    end
+  end
+end

data/lib/llmemory/mcp/server.rb

@@ -0,0 +1,185 @@
+# frozen_string_literal: true
+
+require "mcp"
+require_relative "authentication"
+require_relative "tools/memory_search"
+require_relative "tools/memory_save"
+require_relative "tools/memory_retrieve"
+require_relative "tools/memory_timeline"
+require_relative "tools/memory_add_message"
+require_relative "tools/memory_consolidate"
+require_relative "tools/memory_stats"
+require_relative "tools/memory_info"
+require_relative "tools/memory_timeline_context"
+
+module Llmemory
+  module MCP
+    class Server
+      attr_reader :server
+
+      def initialize(name: "llmemory", version: nil)
+        @server = ::MCP::Server.new(
+          name: name,
+          version: version || Llmemory::VERSION,
+          instructions: instructions_text,
+          tools: all_tools
+        )
+      end
+
+      def run_stdio
+        transport = ::MCP::Server::Transports::StdioTransport.new(@server)
+        transport.open
+      end
+
+      def run_http(port: 3100, host: "0.0.0.0", ssl_cert: nil, ssl_key: nil)
+        require "webrick"
+
+        @http_transport = ::MCP::Server::Transports::StreamableHTTPTransport.new(@server)
+        app = build_rack_app(@http_transport)
+
+        webrick_options = {
+          Port: port,
+          BindAddress: host,
+          Logger: WEBrick::Log.new($stderr, WEBrick::Log::INFO),
+          AccessLog: []
+        }
+
+        # Configure SSL/HTTPS if certificates provided
+        use_ssl = ssl_cert && ssl_key
+        if use_ssl
+          require "webrick/https"
+          require "openssl"
+
+          webrick_options[:SSLEnable] = true
+          webrick_options[:SSLCertificate] = OpenSSL::X509::Certificate.new(File.read(ssl_cert))
+          webrick_options[:SSLPrivateKey] = OpenSSL::PKey::RSA.new(File.read(ssl_key))
+        end
+
+        webrick_server = WEBrick::HTTPServer.new(webrick_options)
+
+        webrick_server.mount_proc "/" do |req, res|
+          rack_env = build_rack_env(req)
+          status, headers, body = app.call(rack_env)
+
+          res.status = status
+          headers.each { |key, value| res[key] = value }
+          res.body = body.is_a?(Array) ? body.join : body
+        end
+
+        trap("INT") { webrick_server.shutdown }
+        trap("TERM") { webrick_server.shutdown }
+
+        protocol = use_ssl ? "https" : "http"
+        $stderr.puts "llmemory MCP server listening on #{protocol}://#{host}:#{port}"
+        $stderr.puts "Authentication: #{ENV["MCP_TOKEN"] ? "enabled (MCP_TOKEN set)" : "disabled"}"
+
+        webrick_server.start
+      ensure
+        @http_transport&.close
+      end
+
+      # Returns a Rack app for use with custom servers (Puma, etc.)
+      def rack_app
+        transport = ::MCP::Server::Transports::StreamableHTTPTransport.new(@server)
+        build_rack_app(transport)
+      end
+
+      private
+
+      def build_rack_app(transport)
+        app = ->(env) { transport.handle_request(RackRequest.new(env)) }
+
+        # Wrap with authentication if MCP_TOKEN is set
+        if ENV["MCP_TOKEN"]
+          app = Authentication.new(app)
+        end
+
+        app
+      end
+
+      def build_rack_env(req)
+        env = {
+          "REQUEST_METHOD" => req.request_method,
+          "SCRIPT_NAME" => "",
+          "PATH_INFO" => req.path,
+          "QUERY_STRING" => req.query_string || "",
+          "SERVER_NAME" => req.host,
+          "SERVER_PORT" => req.port.to_s,
+          "HTTP_HOST" => req["Host"],
+          "rack.input" => StringIO.new(req.body || ""),
+          "rack.url_scheme" => req.ssl? ? "https" : "http"
+        }
+
+        # Copy HTTP headers
+        req.header.each do |key, values|
+          header_key = "HTTP_#{key.upcase.tr("-", "_")}"
+          env[header_key] = values.is_a?(Array) ? values.join(", ") : values
+        end
+
+        env
+      end
+
+      # Simple wrapper to make Rack env look like a Rack request
+      class RackRequest
+        def initialize(env)
+          @env = env
+        end
+
+        def env
+          @env
+        end
+
+        def body
+          @env["rack.input"]
+        end
+
+        def params
+          @params ||= parse_query_string(@env["QUERY_STRING"] || "")
+        end
+
+        private
+
+        def parse_query_string(qs)
+          qs.split("&").each_with_object({}) do |pair, hash|
+            key, value = pair.split("=", 2)
+            hash[key] = value if key
+          end
+        end
+      end
+
+      def all_tools
+        [
+          Tools::MemorySearch,
+          Tools::MemorySave,
+          Tools::MemoryRetrieve,
+          Tools::MemoryTimeline,
+          Tools::MemoryTimelineContext,
+          Tools::MemoryAddMessage,
+          Tools::MemoryConsolidate,
+          Tools::MemoryStats,
+          Tools::MemoryInfo
+        ]
+      end
+
+      def instructions_text
+        <<~INSTRUCTIONS
+          llmemory MCP Server - Persistent Memory System for LLM Agents
+
+          This server provides tools to manage persistent memory for users:
+          - Search memories by semantic query
+          - Save new observations and facts
+          - Retrieve context optimized for LLM inference
+          - Manage conversation sessions
+          - Consolidate conversations into long-term memory
+
+          RECOMMENDED WORKFLOW:
+          1. At start of conversation: use memory_retrieve to get relevant context
+          2. During conversation: use memory_save for important observations
+          3. At end: use memory_consolidate to persist the conversation
+
+          Use memory_info tool for detailed usage guidelines.
+        INSTRUCTIONS
+      end
+    end
+  end
+end

data/lib/llmemory/mcp/tools/memory_add_message.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module MCP
+    module Tools
+      class MemoryAddMessage < ::MCP::Tool
+        description "Add a message to the short-term conversation memory."
+
+        input_schema(
+          properties: {
+            user_id: { type: "string", description: "User identifier" },
+            session_id: { type: "string", description: "Session identifier (default: 'default')" },
+            role: { type: "string", enum: ["user", "assistant", "system"], description: "Message role" },
+            content: { type: "string", description: "Message content" }
+          },
+          required: ["user_id", "role", "content"]
+        )
+
+        class << self
+          def call(user_id:, role:, content:, session_id: nil, server_context: nil)
+            session = session_id || "default"
+
+            memory = Llmemory::Memory.new(user_id: user_id, session_id: session)
+            memory.add_message(role: role.to_sym, content: content)
+
+            ::MCP::Tool::Response.new([{
+              type: "text",
+              text: "Message added to session '#{session}'.\nRole: #{role}\nContent: #{truncate(content, 100)}"
+            }])
+          rescue => e
+            ::MCP::Tool::Response.new([{
+              type: "text",
+              text: "Error adding message: #{e.message}"
+            }], error: true)
+          end
+
+          private
+
+          def truncate(text, max_length)
+            return text if text.length <= max_length
+            "#{text[0, max_length]}..."
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/mcp/tools/memory_consolidate.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module MCP
+    module Tools
+      class MemoryConsolidate < ::MCP::Tool
+        description "Consolidate current short-term conversation into long-term memory. Extracts facts and observations from the conversation."
+
+        input_schema(
+          properties: {
+            user_id: { type: "string", description: "User identifier" },
+            session_id: { type: "string", description: "Session identifier (default: 'default')" },
+            clear_session: { type: "boolean", description: "Clear session after consolidation (default: false)" }
+          },
+          required: ["user_id"]
+        )
+
+        class << self
+          def call(user_id:, session_id: nil, clear_session: nil, server_context: nil)
+            session = session_id || "default"
+            should_clear = clear_session || false
+
+            memory = Llmemory::Memory.new(user_id: user_id, session_id: session)
+
+            messages = memory.messages
+            if messages.empty?
+              return ::MCP::Tool::Response.new([{
+                type: "text",
+                text: "No messages to consolidate in session '#{session}'."
+              }])
+            end
+
+            message_count = messages.size
+            memory.consolidate!
+
+            if should_clear
+              memory.clear_session!
+            end
+
+            response_text = "Consolidated #{message_count} messages from session '#{session}' into long-term memory."
+            response_text += "\nSession cleared." if should_clear
+
+            ::MCP::Tool::Response.new([{
+              type: "text",
+              text: response_text
+            }])
+          rescue => e
+            ::MCP::Tool::Response.new([{
+              type: "text",
+              text: "Error consolidating memory: #{e.message}"
+            }], error: true)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/mcp/tools/memory_info.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module MCP
+    module Tools
+      class MemoryInfo < ::MCP::Tool
+        description "Get documentation on how to use llmemory tools effectively. Read this first to understand the memory system."
+
+        input_schema(
+          properties: {}
+        )
+
+        class << self
+          def call(server_context: nil)
+            ::MCP::Tool::Response.new([{
+              type: "text",
+              text: documentation_text
+            }])
+          end
+
+          private
+
+          def documentation_text
+            <<~DOC
+              # llmemory - Memory System for LLM Agents
+
+              ## Overview
+              llmemory provides persistent memory across conversations, combining:
+              - **Short-term memory**: Recent conversation messages per session
+              - **Long-term memory**: Extracted facts, preferences, and observations
+
+              ## Recommended Workflow
+
+              ### 1. Start of Conversation
+              Use memory_retrieve to get relevant context:
+              ```
+              memory_retrieve(query: "<user's first message>", user_id: "user123")
+              ```
+              This returns relevant context from both short and long-term memory.
+
+              ### 2. During Conversation
+              Save important observations as you learn them:
+              ```
+              memory_save(user_id: "user123", content: "User prefers concise answers")
+              ```
+
+              ### 3. End of Conversation
+              Consolidate the conversation to extract and store facts:
+              ```
+              memory_consolidate(user_id: "user123", session_id: "default")
+              ```
+
+              ## Tools Reference
+
+              | Tool | Purpose |
+              |------|---------|
+              | memory_search | Find specific memories by query |
+              | memory_save | Store a new observation/fact |
+              | memory_retrieve | Get context for LLM inference |
+              | memory_timeline | Get recent memories chronologically |
+              | memory_timeline_context | Get N items before/after a specific memory |
+              | memory_add_message | Add message to short-term |
+              | memory_consolidate | Extract facts from conversation |
+              | memory_stats | Get memory statistics |
+
+              ## Best Practices
+
+              1. **Be specific** when saving observations
+              2. **Use categories** to organize facts (preferences, work, personal, technical)
+              3. **Consolidate regularly** to not lose important information
+              4. **Search before asking** - check if you already know something
+              5. **Use session_id** to separate different conversation contexts
+            DOC
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/mcp/tools/memory_retrieve.rb

@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module MCP
+    module Tools
+      class MemoryRetrieve < ::MCP::Tool
+        description "Retrieve relevant context from memory optimized for LLM inference. Combines short-term conversation history with relevant long-term memories."
+
+        input_schema(
+          properties: {
+            query: { type: "string", description: "Context query (e.g., current user message)" },
+            user_id: { type: "string", description: "User identifier" },
+            session_id: { type: "string", description: "Session identifier (default: 'default')" },
+            max_tokens: { type: "integer", description: "Maximum tokens for context (default: 2000)" },
+            include_timeline_context: { type: "boolean", description: "Include N items before/after top matches (default: false)" },
+            timeline_window: { type: "integer", description: "Number of items before/after for timeline context (default: 3)" }
+          },
+          required: ["query", "user_id"]
+        )
+
+        class << self
+          def call(query:, user_id:, session_id: nil, max_tokens: nil, include_timeline_context: nil, timeline_window: nil, server_context: nil)
+            session = session_id || "default"
+            tokens = max_tokens || 2000
+            include_timeline = include_timeline_context == true
+            window = timeline_window || 3
+
+            memory = Llmemory::Memory.new(user_id: user_id, session_id: session)
+            context = memory.retrieve(query, max_tokens: tokens)
+
+            # Add timeline context if requested
+            if include_timeline && !context.to_s.strip.empty?
+              timeline_context = fetch_timeline_context(user_id, query, window)
+              context = "#{context}\n\n#{timeline_context}" unless timeline_context.empty?
+            end
+
+            if context.to_s.strip.empty?
+              ::MCP::Tool::Response.new([{
+                type: "text",
+                text: "No relevant context found for this query."
+              }])
+            else
+              ::MCP::Tool::Response.new([{
+                type: "text",
+                text: context
+              }])
+            end
+          rescue => e
+            ::MCP::Tool::Response.new([{
+              type: "text",
+              text: "Error retrieving context: #{e.message}"
+            }], error: true)
+          end
+
+          private
+
+          def fetch_timeline_context(user_id, query, window)
+            storage = build_storage
+            items = storage.search_items(user_id, query)
+            return "" if items.empty?
+
+            # Get timeline context around the first match
+            top_item = items.first
+            item_id = top_item[:id] || top_item["id"]
+            return "" unless item_id
+
+            result = storage.get_items_around(user_id, item_id, before: window, after: window)
+            format_timeline_context(result, item_id)
+          rescue
+            ""
+          end
+
+          def build_storage
+            if Llmemory.configuration.long_term_type.to_s == "graph_based"
+              LongTerm::GraphBased::Storages.build
+            else
+              LongTerm::FileBased::Storages.build
+            end
+          end
+
+          def format_timeline_context(result, reference)
+            return "" if result[:before].empty? && result[:target].nil? && result[:after].empty?
+
+            lines = ["=== TIMELINE CONTEXT ===", ""]
+            lines << "Events around match '#{reference}':"
+            lines << ""
+
+            result[:before].each do |item|
+              lines << "  [BEFORE] #{format_item(item)}"
+            end
+
+            if result[:target]
+              lines << "  [MATCH]  #{format_item(result[:target])}"
+            end
+
+            result[:after].each do |item|
+              lines << "  [AFTER]  #{format_item(item)}"
+            end
+
+            lines << ""
+            lines << "=== END TIMELINE CONTEXT ==="
+            lines.join("\n")
+          end
+
+          def format_item(item)
+            ts = item[:created_at] || item["created_at"]
+            ts_str = ts.respond_to?(:strftime) ? ts.strftime("%Y-%m-%d") : ts.to_s[0, 10]
+            content = item[:content] || item["content"]
+            category = item[:category] || item["category"]
+            cat_str = category ? "[#{category}] " : ""
+            "#{ts_str} #{cat_str}#{truncate(content, 80)}"
+          end
+
+          def truncate(text, max)
+            return text.to_s if text.to_s.length <= max
+            "#{text.to_s[0, max]}..."
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/mcp/tools/memory_save.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module MCP
+    module Tools
+      class MemorySave < ::MCP::Tool
+        description "Save a new observation or fact to long-term memory. Use this to remember important information about the user."
+
+        input_schema(
+          properties: {
+            user_id: { type: "string", description: "User identifier" },
+            content: { type: "string", description: "The observation or fact to remember" },
+            category: {
+              type: "string",
+              description: "Category for the memory (e.g., preferences, work, personal). If omitted, will be auto-classified."
+            }
+          },
+          required: ["user_id", "content"]
+        )
+
+        class << self
+          def call(user_id:, content:, category: nil, server_context: nil)
+            storage = build_storage
+
+            # If no category provided, use a default
+            cat = category || "observations"
+
+            # Generate a simple resource ID for tracking
+            resource_id = "mcp_#{Time.now.to_i}_#{rand(1000)}"
+
+            storage.save_item(
+              user_id,
+              category: cat,
+              content: content,
+              source_resource_id: resource_id
+            )
+
+            ::MCP::Tool::Response.new([{
+              type: "text",
+              text: "Memory saved successfully.\nCategory: #{cat}\nContent: #{content}"
+            }])
+          rescue => e
+            ::MCP::Tool::Response.new([{
+              type: "text",
+              text: "Error saving memory: #{e.message}"
+            }], error: true)
+          end
+
+          private
+
+          def build_storage
+            LongTerm::FileBased::Storages.build
+          end
+        end
+      end
+    end
+  end
+end