diff --git a/lib/phoenix_live_session.ex b/lib/phoenix_live_session.ex
index 10f5214..4f75202 100644
--- a/lib/phoenix_live_session.ex
+++ b/lib/phoenix_live_session.ex
@@ -225,13 +225,11 @@ defmodule PhoenixLiveSession do
   """
   @spec maybe_subscribe(Phoenix.LiveView.Socket.t(), Plug.Session.Store.session()) ::
           Phoenix.LiveView.Socket.t()
-  def maybe_subscribe(socket, session) do
+  def maybe_subscribe(socket, %{"__sid__" => sid, "__opts__" => opts}) do
     if LiveView.connected?(socket) do
-      sid = Map.fetch!(session, :__sid__)
-      opts = Map.fetch!(session, :__opts__)
       pub_sub = Keyword.fetch!(opts, :pub_sub)
-      channel = "live_session:#{sid}"
-      PubSub.subscribe(pub_sub, channel)
+
+      PubSub.subscribe(pub_sub, "live_session:#{sid}")
 
       put_in(socket.private[:live_session], id: sid, opts: opts)
     else
@@ -239,6 +237,8 @@ defmodule PhoenixLiveSession do
     end
   end
 
+  def maybe_subscribe(socket, _), do: socket
+
   @doc """
   This function can be called in two ways:any()
 
@@ -253,8 +253,8 @@ defmodule PhoenixLiveSession do
   from the `mount/3` callback directly in this function.
   Retrieves and returns updated session data.
   """
-  @spec put_session(Phoenix.LiveView.Socket.t(), String.t() | atom(), term()) ::
-          Phoenix.LiveView.Socket.t()
+  @spec put_session(Phoenix.LiveView.Socket.t() | map(), String.t() | atom(), term()) ::
+          Phoenix.LiveView.Socket.t() | %{}
   def put_session(%Phoenix.LiveView.Socket{} = socket, key, value) do
     sid = get_in(socket.private, [:live_session, :id])
     opts = get_in(socket.private, [:live_session, :opts])
@@ -263,8 +263,7 @@ defmodule PhoenixLiveSession do
     socket
   end
 
-  @spec put_session(%{__sid__: String.t(), __opts__: list()}, String.t() | atom(), term()) :: %{}
-  def put_session(%{__sid__: sid, __opts__: opts}, key, value) do
+  def put_session(%{"__sid__" => sid, "__opts__" => opts}, key, value) do
     put_in(sid, to_string(key), value, opts)
 
     get(nil, sid, opts)