Don't use private container

Sources/Flash/Extensions/Container+Flash.swift

@@ -0,0 +1,16 @@
+import Vapor
+
+extension Container {
+    /// Add a flash message to the `Response` of this `Future`.
+    ///
+    /// - Parameters:
+    ///   - type: The type of the flash message (e.g. `.success` or `.error`).
+    ///   - message: The message to display.
+    /// - Returns: The `Response` with the flash message added.
+    public func flash(_ type: Flash.Kind, _ message: String) -> Self {
+        if let container: FlashContainer = try? make() {
+            container.flashes.append(.init(type, message))
+        }
+        return self
+    }
+}

Sources/Flash/Extensions/Future+Flash.swift

@@ -1,9 +1,15 @@
 import Vapor
 
-public extension Future where T: Response {
+public extension Future where T == Response {
+    /// Add a flash message to the `Response` of this `Future`.
+    ///
+    /// - Parameters:
+    ///   - type: The type of the flash message (e.g. `.success` or `.error`).
+    ///   - message: The message to display.
+    /// - Returns: A `Future` containing the `Response` with the added flash message.
     public func flash(_ type: Flash.Kind, _ message: String) -> Future<Response> {
-        return self.map(to: Response.self) { res in
-            return res.flash(type, message)
+        return map { response in
+            response.flash(type, message)
         }
     }
 }

Sources/Flash/Middlewares/FlashMiddleware.swift

@@ -1,41 +1,49 @@
 import Vapor
 
+/// Middleware that:
+/// - decodes flashes from the `Session` and puts them on the `Request`s container.
+/// - encodes flashes from the `Response` to the `Session`.
 public struct FlashMiddleware: Middleware, ServiceType {
-    private static let sessionKey = "_flash"
-
+    /// See `ServiceType.makeService`.
     public static func makeService(for container: Container) throws -> FlashMiddleware {
         return .init()
     }
 
+    /// Create a new `FlashMiddleware`.
     public init() {}
 
     /// See Middleware.respond
     public func respond(to req: Request, chainingTo next: Responder) throws -> Future<Response> {
-        try FlashMiddleware.handle(req: req)
-        return try next.respond(to: req)
-            .map(to: Response.self) { resp in
-                try FlashMiddleware.handle(req: req, resp: resp)
-                return resp
+        try req.decodeFlashDataFromSession()
+        return try next
+            .respond(to: req)
+            .try { resp in
+                try req.encodeFlashDataToSession(from: resp)
             }
     }
+}
 
-    public static func handle(req: Request) throws {
-        let session = try req.session()
+private let sessionKey = "_flash"
 
-        if let data = session[sessionKey]?.data(using: .utf8) {
-            let flash = try JSONDecoder().decode(FlashContainer.self, from: data)
-            let container = try req.privateContainer.make(FlashContainer.self)
-            container.new = flash.new
-            container.old = flash.old
+extension Request {
+    func decodeFlashDataFromSession() throws {
+        let container = try privateContainer.make(FlashContainer.self)
+
+        guard let data = try session()[sessionKey]?.data(using: .utf8) else {
+            container.flashes = []
+            return
         }
+
+        let decoded = try JSONDecoder().decode(FlashContainer.self, from: data)
+        container.flashes = decoded.flashes
     }
 
-    public static func handle(req: Request, resp: Response) throws {
+    func encodeFlashDataToSession(from resp: Response) throws {
         let container = try resp.privateContainer.make(FlashContainer.self)
         let flash = try String(
             data: JSONEncoder().encode(container),
             encoding: .utf8
         )
-        try req.session()[sessionKey] = flash
+        try session()[sessionKey] = flash
     }
 }

Sources/Flash/Models/Flash.swift

@@ -1,21 +1,26 @@
 import Vapor
 
-public final class Flash: Codable {
+/// A message that can be displayed on a web page. The message is displayed
+public struct Flash: Codable {
     public enum Kind: String, Codable {
         case error
         case success
         case info
         case warning
     }
 
-    public var kind: Kind
-    public var message: String
+    /// The kind of message.
+    /// - See: `Flash.Kind`.
+    public let kind: Kind
 
-    public init(kind: Kind, message: String) {
-        self.kind = kind
-        self.message = message
-    }
+    /// The message to be displayed.
+    public let message: String
 
+    /// Create a new `Flash` message.
+    ///
+    /// - Parameters:
+    ///   - kind: The kind of message.
+    ///   - message: The text to be displayed.
     public init(_ kind: Kind, _ message: String) {
         self.kind = kind
         self.message = message

Sources/Flash/Models/FlashContainer.swift

@@ -1,16 +1,6 @@
 import Vapor
 
+/// Container that contains all `Flash`es.
 public final class FlashContainer: Codable, Service {
-    public var new: [Flash] = []
-    public var old: [Flash] = []
-
-    public var flashes: [Flash] {
-        get {
-            return new
-        }
-
-        set {
-            new = newValue
-        }
-    }
+    public var flashes: [Flash] = []
 }

Sources/Flash/Providers/FlashProvider.swift

@@ -2,9 +2,13 @@ import Leaf
 import Sugar
 import Vapor
 
+/// Provider that registers FlashMiddleware and FlashContainer, and FlashTag.
 public final class FlashProvider: Provider {
+
+    /// Create a new `FlashProvider`.
     public init() {}
 
+    /// See `Provider.register`.
     public func register(_ services: inout Services) throws {
         try services.register(MutableLeafTagConfigProvider())
         services.register(FlashMiddleware.self)
@@ -13,6 +17,7 @@ public final class FlashProvider: Provider {
         }
     }
 
+    /// See `Provider.didBoot`
     public func didBoot(_ container: Container) throws -> EventLoopFuture<Void> {
         let tags: MutableLeafTagConfig = try container.make()
         tags.use(FlashTag(), as: "flash")

Sources/Flash/Tags/FlashTag.swift

@@ -1,70 +1,63 @@
 import Leaf
 import TemplateKit
 
+/// Tag that
 public final class FlashTag: TagRenderer {
-    public func render(tag: TagContext) throws -> EventLoopFuture<TemplateData> {
+    /// Create a new `FlashTag`.
+    public init() {}
+
+    /// See `TagRenderer.render`.
+    public func render(tag: TagContext) throws -> Future<TemplateData> {
         let body = try tag.requireBody()
-        let flash = try tag.container.make(FlashContainer.self)
+        let flash: FlashContainer = try tag.container.make()
 
         guard !flash.flashes.isEmpty else {
-           return Future.map(on: tag) {
-                .string("")
-            }
+            return tag.future(.null)
         }
 
-        var dict = tag.context.data.dictionary ?? [:]
-        dict["all"] = try .array(flash.flashes.map {
-            try $0.convertToTemplateData()
-        })
-
-        dict["errors"] = try .array(flash.flashes.compactMap { flash in
-            guard flash.kind == .error else { return nil }
-            return try flash.convertToTemplateData()
-        })
-
-        dict["warnings"] = try .array(flash.flashes.compactMap { flash in
-            guard flash.kind == .warning else { return nil }
-            return try flash.convertToTemplateData()
-        })
-
-        dict["successes"] = try .array(flash.flashes.compactMap { flash in
-            guard flash.kind == .success else { return nil }
-            return try flash.convertToTemplateData()
-        })
+        var flashes =
+            try Dictionary(grouping: flash.flashes) { flash in
+                flash.kind.groupingKey
+            }
+            .mapValues { flashes -> TemplateData in
+                .array(try flashes.map { flash in
+                    try flash.convertToTemplateData()
+                })
+            }
 
-        dict["information"] = try .array(flash.flashes.compactMap { flash in
-            guard flash.kind == .info else { return nil }
-            return try flash.convertToTemplateData()
-        })
+        flashes["all"] = .array(flashes.values.compactMap { $0 })
 
-        tag.context.data = .dictionary(dict)
+        let existing = (tag.context.data.dictionary ?? [:])
+        tag.context.data = .dictionary(
+            existing.merging(flashes) { (existing, fromFlash) -> TemplateData in
+                return fromFlash
+            }
+        )
 
-        return tag.serializer.serialize(ast: body).map(to: TemplateData.self) { er in
-            let body = String(data: er.data, encoding: .utf8) ?? ""
-            return .string(body)
+        return tag.serializer.serialize(ast: body).map { view in
+            .string(String(data: view.data, encoding: .utf8) ?? "")
         }
     }
-
-    public init() {}
 }
 
 extension Flash: TemplateDataRepresentable {
+    /// See `TemplateDataRepresentable`.
     public func convertToTemplateData() throws -> TemplateData {
-        return TemplateData.dictionary([
-            "kind": .string(self.kind.rawValue),
-            "bootstrapClass": .string(self.kind.bootstrapClass),
-            "message": .string(self.message)
+        return .dictionary([
+            "kind": .string(kind.rawValue),
+            "bootstrapClass": .string(kind.rawValue),
+            "message": .string(message)
         ])
     }
 }
 
 extension Flash.Kind {
-    var bootstrapClass: String {
+    var groupingKey: String {
         switch self {
-        case .error: return "danger"
-        case .warning: return "warning"
-        case .success: return "success"
-        case .info: return "info"
+        case .error: return "errors"
+        case .info: return "information"
+        case .success: return "successes"
+        case .warning: return "warnings"
         }
     }
 }