Stream IMAP message parts instead of buffering whole messages

lib/mail_mcp.rb

@@ -16,6 +16,9 @@ class << self
 
 require_relative "mail_mcp/jwt_service"
 require_relative "mail_mcp/pkce"
+require_relative "mail_mcp/base64_stream"
+require_relative "mail_mcp/message_structure"
+require_relative "mail_mcp/message_reader"
 require_relative "mail_mcp/imap_client"
 require_relative "mail_mcp/smtp_client"
 require_relative "mail_mcp/attachment_store"

lib/mail_mcp/attachment_store.rb

@@ -5,28 +5,43 @@ module MailMCP
   module AttachmentStore
     EXPIRY = 7 * 24 * 3600
 
-    def self.upload(content:, filename:, content_type:)
-      key = "attachments/#{SecureRandom.uuid}/#{filename}"
-      bucket = ENV.fetch("AWS_S3_BUCKET")
-
-      s3.put_object(
-        bucket: bucket,
-        key: key,
-        body: content,
-        content_type: content_type
-      )
-
-      presigner.presigned_url(:get_object, bucket: bucket, key: key, expires_in: EXPIRY)
-    end
+    class << self
+      # Uploads from an open IO (typically a Tempfile), letting the SDK read it in
+      # chunks instead of holding the whole attachment as a String.
+      def upload_io(io:, filename:, content_type:)
+        io.rewind
+        key = "attachments/#{SecureRandom.uuid}/#{filename}"
+        bucket = ENV.fetch("AWS_S3_BUCKET")
 
-    def self.s3
-      Aws::S3::Client.new
-    end
-    private_class_method :s3
+        s3.put_object(
+          bucket: bucket,
+          key: key,
+          body: io,
+          content_type: content_type
+        )
+
+        presigner.presigned_url(:get_object, bucket: bucket, key: key, expires_in: EXPIRY)
+      end
+
+      # Drops the memoized clients. Only needed by tests and after credential changes.
+      def reset!
+        @s3 = nil
+        @presigner = nil
+      end
+
+      private
+
+      # Memoized because the first Aws::S3::Client.new costs ~110ms and ~19MB of RSS
+      # loading plugins, and the un-memoized version built two per attachment — once
+      # here and once more inside #presigner. Aws::S3::Client is thread-safe, so one
+      # instance serves every Puma thread.
+      def s3
+        @s3 ||= Aws::S3::Client.new
+      end
 
-    def self.presigner
-      Aws::S3::Presigner.new(client: s3)
+      def presigner
+        @presigner ||= Aws::S3::Presigner.new(client: s3)
+      end
     end
-    private_class_method :presigner
   end
 end

lib/mail_mcp/base64_stream.rb

@@ -0,0 +1,32 @@
+module MailMCP
+  # Incremental base64 decoder. Feed it arbitrary slices of a base64 stream and it
+  # emits decoded bytes as soon as it holds complete four-character groups, carrying
+  # the remainder (always fewer than four characters) into the next call. This is what
+  # lets an attachment of any size be decoded with a bounded amount of memory.
+  class Base64Stream
+    GROUP = 4
+
+    def initialize
+      @carry = +""
+    end
+
+    # Consumes +chunk+: it is stripped and sliced in place so that decoding a 4MiB
+    # chunk does not allocate three more copies of it. Callers must read anything they
+    # need from the chunk (its size, in particular) before calling this.
+    def push(chunk)
+      chunk = chunk.dup if chunk.frozen?
+      chunk.delete!("\r\n\t ")
+      chunk.prepend(@carry) unless @carry.empty?
+      leftover = chunk.bytesize % GROUP
+      @carry = leftover.zero? ? +"" : chunk.slice!(chunk.bytesize - leftover, leftover)
+      chunk.empty? ? "".b : chunk.unpack1("m")
+    end
+
+    # Decodes whatever is left over, including any padding.
+    def finish
+      remainder = @carry
+      @carry = +""
+      remainder.unpack1("m")
+    end
+  end
+end

lib/mail_mcp/imap_client.rb

@@ -75,16 +75,14 @@ def list_messages(folder:, page: 1, per_page: 20)
       { messages: messages, total: total, page: page, per_page: per_page }
     end
 
+    # Delegates to MessageReader, which fetches the message a part at a time so peak
+    # memory tracks the largest part rather than the whole message.
     def get_message(folder:, uid:)
       MailMCP.logger.info { "IMAP get_message folder=#{folder.inspect} uid=#{uid}" }
       @imap.examine(folder)
-      data = @imap.uid_fetch([uid.to_i], %w[RFC822 FLAGS]).first
-      unless data
-        MailMCP.logger.warn { "IMAP get_message not found folder=#{folder.inspect} uid=#{uid}" }
-        return nil
-      end
-
-      format_message(uid: uid, parsed: Mail.new(data.attr["RFC822"]), flags: data.attr["FLAGS"])
+      message = MessageReader.new(@imap, uid.to_i).read
+      MailMCP.logger.warn { "IMAP get_message not found folder=#{folder.inspect} uid=#{uid}" } unless message
+      message
     end
 
     def search_messages(folder:, query:)
@@ -142,27 +140,6 @@ def self.open_connection(config)
 
     private
 
-    def format_message(uid:, parsed:, flags:)
-      {
-        uid: uid,
-        message_id: parsed.message_id,
-        in_reply_to: parsed.in_reply_to,
-        references: parsed.references,
-        subject: parsed.subject,
-        from: parsed.from,
-        sender: parsed.sender,
-        reply_to: parsed.reply_to,
-        to: parsed.to,
-        cc: parsed.cc,
-        bcc: parsed.bcc,
-        date: parsed.date&.iso8601,
-        text_body: parsed.text_part&.decoded,
-        html_body: parsed.html_part&.decoded,
-        flags: flags,
-        attachments: extract_attachments(parsed)
-      }
-    end
-
     def format_envelope(msg)
       env = msg.attr["ENVELOPE"]
       {
@@ -187,19 +164,5 @@ def format_addresses(addrs)
 
       addrs.map { |a| "#{a.name} <#{a.mailbox}@#{a.host}>" }
     end
-
-    def extract_attachments(mail)
-      mail.attachments.map do |att|
-        # Mail#decoded does not memoize — it base64-decodes into a new String on
-        # every call, so decode once and reuse for both the upload and the size.
-        content = att.decoded
-        url = AttachmentStore.upload(
-          content: content,
-          filename: att.filename || "attachment",
-          content_type: att.content_type
-        )
-        { filename: att.filename, content_type: att.content_type, size: content.bytesize, url: url }
-      end
-    end
   end
 end

lib/mail_mcp/message_reader.rb

@@ -0,0 +1,161 @@
+require "tempfile"
+
+module MailMCP
+  # Reads one message over an existing IMAP connection a part at a time, so peak memory
+  # tracks the part being handled rather than the size of the whole message. Fetching
+  # RFC822 instead would materialize the entire message as a single String.
+  class MessageReader
+    # Raised when the server answers a section fetch with a section we did not ask for.
+    # Treated as an error because the alternative is a silently truncated attachment.
+    class IncompletePart < StandardError; end
+
+    # Encoded bytes pulled per FETCH when streaming an attachment. Bounds peak memory
+    # regardless of attachment size; larger chunks mean fewer round trips but hold a
+    # Puma thread for longer.
+    CHUNK_SIZE = 4 * 1024 * 1024
+
+    def initialize(imap, uid)
+      @imap = imap
+      @uid = uid
+    end
+
+    # Returns the message as the tool-facing hash, or nil when the uid is gone.
+    def read
+      data = @imap.uid_fetch([@uid], ["BODYSTRUCTURE", "FLAGS", "BODY.PEEK[HEADER]"])&.first
+      return nil unless data
+
+      parts = MessageStructure.flatten(data.attr["BODYSTRUCTURE"])
+      MailMCP.logger.debug { "IMAP message uid=#{@uid} parts=#{parts.map(&:section).inspect}" }
+
+      to_h(headers: Mail.new(data.header.to_s), flags: data.attr["FLAGS"], parts: parts)
+    end
+
+    private
+
+    def to_h(headers:, flags:, parts:)
+      inline = parts.reject(&:attachment?)
+      {
+        uid: @uid,
+        message_id: headers.message_id,
+        in_reply_to: headers.in_reply_to,
+        references: headers.references,
+        subject: headers.subject,
+        from: headers.from,
+        sender: headers.sender,
+        reply_to: headers.reply_to,
+        to: headers.to,
+        cc: headers.cc,
+        bcc: headers.bcc,
+        date: headers.date&.iso8601,
+        text_body: inline_body(inline, "text/plain"),
+        html_body: inline_body(inline, "text/html"),
+        flags: flags,
+        attachments: upload_attachments(parts.select(&:attachment?))
+      }
+    end
+
+    # Mail#text_part/#html_part pick the first non-attachment part of that type, so
+    # match the flattened list the same way.
+    def inline_body(inline, mime_type)
+      part = inline.find { |candidate| candidate.mime_type == mime_type }
+      return nil unless part
+
+      body = fetch_section(part.section)
+      body && decode_inline(part, body)
+    end
+
+    # Letting Mail decode a rebuilt part applies both the transfer-encoding and the
+    # charset conversion, which is what #decoded did on the old whole-message parse.
+    # Hand-rolling that would mislabel any non-UTF-8 body.
+    def decode_inline(part, body)
+      rebuilt = Mail::Part.new
+      rebuilt.content_type = part.charset ? "#{part.mime_type}; charset=#{part.charset}" : part.mime_type
+      rebuilt.content_transfer_encoding = part.encoding
+      rebuilt.body = body
+      rebuilt.decoded
+    rescue StandardError => e
+      MailMCP.logger.warn { "IMAP inline decode failed section=#{part.section}: #{e.class}: #{e.message}" }
+      nil
+    end
+
+    # Each attachment is streamed to disk, uploaded, then discarded before the next is
+    # fetched, so attachments never accumulate in memory.
+    def upload_attachments(parts)
+      parts.map do |part|
+        Tempfile.create("mail_mcp_attachment", binmode: true) do |io|
+          size = fetch_attachment(io, part)
+          url = AttachmentStore.upload_io(
+            io: io,
+            filename: part.filename || "attachment",
+            content_type: part.mime_type
+          )
+          { filename: part.filename, content_type: part.mime_type, size: size, url: url }
+        end
+      end
+    end
+
+    # Returns the decoded byte count written to +io+.
+    def fetch_attachment(io, part)
+      return stream_section(io, part) if part.streamable?
+
+      # Quoted-printable has to be decoded in one pass, so this path is bounded by the
+      # part size rather than by CHUNK_SIZE. Logged because it is the one place where
+      # the memory guarantee above does not hold.
+      MailMCP.logger.warn do
+        "IMAP attachment decoded in one pass encoding=#{part.encoding.inspect} " \
+          "encoded_size=#{part.encoded_size} section=#{part.section}"
+      end
+      io.write(decode_whole(require_section(part.section), part.encoding))
+    end
+
+    def stream_section(io, part)
+      decoder = part.base64? ? Base64Stream.new : nil
+      offset = 0
+      written = 0
+      loop do
+        fetched, wrote = consume_chunk(io, decoder, part, offset)
+        break if fetched.zero?
+
+        offset += fetched
+        written += wrote
+      end
+      written += io.write(decoder.finish) if decoder
+      written
+    end
+
+    # Kept in its own method so the chunk becomes unreachable as soon as it returns.
+    # Held in a local in the loop above, a CHUNK_SIZE buffer would stay alive across
+    # the next fetch, doubling the resident cost of streaming.
+    # Returns [encoded bytes fetched, decoded bytes written].
+    def consume_chunk(io, decoder, part, offset)
+      chunk = require_section(part.section, offset: offset)
+      return [0, 0] if chunk.empty?
+
+      # #push consumes the chunk, so read its size before handing it over.
+      fetched = chunk.bytesize
+      [fetched, io.write(decoder ? decoder.push(chunk) : chunk)]
+    end
+
+    def decode_whole(raw, encoding)
+      encoder = Mail::Encodings.get_encoding(encoding)
+      encoder ? encoder.decode(raw) : raw
+    rescue StandardError => e
+      MailMCP.logger.warn { "IMAP attachment decode failed encoding=#{encoding.inspect}: #{e.class}: #{e.message}" }
+      raw
+    end
+
+    def require_section(section, offset: nil)
+      fetch_section(section, offset: offset) ||
+        raise(IncompletePart, "IMAP returned no BODY[#{section}] at offset #{offset.inspect} for uid #{@uid}")
+    end
+
+    # net-imap reconciles the request spec against the response key, which differ:
+    # servers answer BODY.PEEK[...] with BODY[...] and append the origin octet on a
+    # partial fetch. A nil result means a section we did not ask for.
+    def fetch_section(section, offset: nil)
+      spec = offset ? "BODY.PEEK[#{section}]<#{offset}.#{CHUNK_SIZE}>" : "BODY.PEEK[#{section}]"
+      data = @imap.uid_fetch([@uid], [spec])&.first
+      data&.part(*section.split("."), offset: offset)
+    end
+  end
+end