nri-prelude: thread session id through LogHandler

nri-prelude/src/Platform.hs

@@ -12,6 +12,8 @@ module Platform
     Internal.LogHandler,
     logHandler,
     requestId,
+    sessionId,
+    setSessionIdIO,
     silentHandler,
     Internal.silentTrack,
 
@@ -64,6 +66,7 @@ where
 import Basics
 import qualified Control.Exception.Safe as Exception
 import qualified Control.Monad.Catch as Catch
+import qualified Data.IORef as IORef
 import qualified Data.Text
 import qualified GHC.Stack as Stack
 import NriPrelude
@@ -170,6 +173,26 @@ logHandler = Internal.Task (pure << Ok)
 requestId :: Task e Text
 requestId = map Internal.requestId logHandler
 
+-- | Read the current request's session id, set by `setSessionIdIO` at
+-- the request boundary (typically a WAI middleware on the application
+-- root). `Platform.Analytics.Internal.trackEvent` uses this to stamp
+-- `session_id` onto every analytics event automatically.
+sessionId :: Task e (Maybe Text)
+sessionId =
+  Internal.Task
+    ( \handler -> do
+        mSid <- IORef.readIORef (Internal.sessionIdRef handler)
+        pure (Ok mSid)
+    )
+
+-- | Set the current request's session id. Intended to be called once at
+-- the request boundary from IO (typically a WAI middleware that has
+-- access to the per-request `LogHandler`). All descendants of the same
+-- request share the underlying ref, so a single write is visible to
+-- every child handler and tracing span.
+setSessionIdIO :: Internal.LogHandler -> Maybe Text -> IO ()
+setSessionIdIO handler = IORef.writeIORef (Internal.sessionIdRef handler)
+
 -- | A log handler that doesn't log anything.
 silentHandler :: IO Internal.LogHandler
 silentHandler = pure Internal.nullHandler

nri-prelude/src/Platform/Analytics/Internal.hs

@@ -14,25 +14,41 @@ module Platform.Analytics.Internal
 where
 
 import qualified Data.Aeson as Aeson
+import qualified Data.Aeson.KeyMap as KeyMap
 import NriPrelude
 import qualified Platform
 import qualified Platform.Internal as Internal
 import Task (Task)
 import qualified Prelude
 
 -- | Send an analytics event. Opens a child tracing span named
--- @analytics.track@, attaches the JSON payload as the span's details,
--- and synchronously invokes the `LogHandler`'s analytics callback.
+-- @analytics.track@, stamps the current request's session id (if any)
+-- onto the event, attaches the resulting JSON payload as the span's
+-- details, and synchronously invokes the `LogHandler`'s analytics
+-- callback.
 trackEvent :: (Aeson.ToJSON e) => e -> Task err ()
 trackEvent event =
-  let value = Aeson.toJSON event
-   in Platform.tracingSpan "analytics.track" <| do
-        Platform.setTracingSpanDetails (AnalyticsEventDetails value)
-        Internal.Task
-          ( \handler -> do
-              Internal.trackAnalyticsEventIO handler value
-              Prelude.pure (Ok ())
-          )
+  Platform.tracingSpan "analytics.track" <| do
+    mSid <- Platform.sessionId
+    let value = stampSessionId mSid (Aeson.toJSON event)
+    Platform.setTracingSpanDetails (AnalyticsEventDetails value)
+    Internal.Task
+      ( \handler -> do
+          Internal.trackAnalyticsEventIO handler value
+          Prelude.pure (Ok ())
+      )
+
+-- | Shallow-merge `session_id` onto an event payload. We only stamp
+-- when the inbound value is a JSON object — every Event sum-type
+-- variant serializes to an object, so non-object inputs are
+-- defensively left untouched. Existing `session_id` keys on the body
+-- are overwritten so the request-scoped id always wins.
+stampSessionId :: Maybe Text -> Aeson.Value -> Aeson.Value
+stampSessionId Nothing v = v
+stampSessionId (Just sid) v = case v of
+  Aeson.Object body ->
+    Aeson.Object (KeyMap.insert "session_id" (Aeson.String sid) body)
+  other -> other
 
 -- | A `TracingSpanDetails` wrapper around the analytics event payload, so
 -- that the JSON we send to the analytics backend is also attached to the

nri-prelude/src/Platform/Internal.hs

@@ -14,6 +14,7 @@ import qualified Data.Aeson as Aeson
 import qualified Data.Aeson.Encoding as Aeson.Encoding
 import qualified Data.IORef as IORef
 import qualified Data.Text
+import qualified System.IO.Unsafe
 import qualified Data.Typeable as Typeable
 import qualified GHC.Clock as Clock
 import GHC.Generics (Generic)
@@ -579,14 +580,32 @@ data LogHandler = LogHandler
     -- threads this opaque callback from `rootTracingSpanIO` through
     -- every child `LogHandler`. See `Platform.Analytics.Internal.trackEvent`
     -- for the user-facing wrapper. Default: `silentTrack`.
-    trackAnalyticsEventIO :: Aeson.Value -> IO ()
+    trackAnalyticsEventIO :: Aeson.Value -> IO (),
+    -- | Per-request session id, propagated to every child `LogHandler`
+    -- via shared `IORef`. Set once at request boundaries (typically a
+    -- WAI middleware on the application root) by writing through this
+    -- ref. `Platform.Analytics.Internal.trackEvent` reads the ref and
+    -- shallow-merges `session_id` onto each event payload, so callers
+    -- of `Analytics.track` don't have to thread the session id through
+    -- explicitly. Default: a per-handler ref containing `Nothing`.
+    sessionIdRef :: IORef.IORef (Maybe Text)
   }
 
 -- | A no-op analytics callback. Used as the default for `nullHandler`
 -- and for platforms that have not opted in to analytics tracking yet.
 silentTrack :: Aeson.Value -> IO ()
 silentTrack _ = pure ()
 
+-- | A shared sentinel `IORef` used by `nullHandler`. We need the field on
+-- `LogHandler` to be inhabited even in the null case, but `nullHandler`
+-- itself is a pure value, so we mint the ref once at module load via
+-- `unsafePerformIO`. Writes to it are no-ops in spirit (nothing reads
+-- back from a null handler in well-behaved code) but technically
+-- visible if multiple null handlers shared this ref.
+nullSessionIdRef :: IORef.IORef (Maybe Text)
+nullSessionIdRef = System.IO.Unsafe.unsafePerformIO (IORef.newIORef Nothing)
+{-# NOINLINE nullSessionIdRef #-}
+
 -- | Helper that creates one of the handler's above. This is intended for
 -- internal use in this library only and not for exposing. Outside of this
 -- library the @rootTracingSpanIO@ is the more user-friendly way to get hands
@@ -597,13 +616,16 @@ mkHandler ::
   Clock ->
   -- | Analytics callback, propagated to every descendant `LogHandler`.
   (Aeson.Value -> IO ()) ->
+  -- | Session-id ref, shared with every descendant `LogHandler` so a
+  -- write at the request root is visible to all children.
+  IORef.IORef (Maybe Text) ->
   -- Finalizer for this loghandler
   (TracingSpan -> IO ()) ->
   -- Root finalizer
   Maybe (TracingSpan -> IO ()) ->
   Text ->
   IO LogHandler
-mkHandler requestId clock trackEventIO onFinish onFinishRoot' name' = do
+mkHandler requestId clock trackEventIO sidRef onFinish onFinishRoot' name' = do
   let onFinishRoot = Maybe.withDefault onFinish onFinishRoot'
   tracingSpanRef <-
     Stack.withFrozenCallStack startTracingSpan clock name'
@@ -612,8 +634,8 @@ mkHandler requestId clock trackEventIO onFinish onFinishRoot' name' = do
   pure
     LogHandler
       { requestId,
-        startChildTracingSpan = mkHandler requestId clock trackEventIO (appendTracingSpanToParent tracingSpanRef) (Just onFinishRoot),
-        startNewRoot = mkHandler requestId clock trackEventIO onFinishRoot Nothing,
+        startChildTracingSpan = mkHandler requestId clock trackEventIO sidRef (appendTracingSpanToParent tracingSpanRef) (Just onFinishRoot),
+        startNewRoot = mkHandler requestId clock trackEventIO sidRef onFinishRoot Nothing,
         setTracingSpanDetailsIO = \details' ->
           updateIORef
             tracingSpanRef
@@ -627,7 +649,8 @@ mkHandler requestId clock trackEventIO onFinish onFinishRoot' name' = do
             tracingSpanRef
             (\tracingSpan' -> tracingSpan' {succeeded = succeeded tracingSpan' ++ Failed, containsFailures = True}),
         finishTracingSpan = finalizeTracingSpan clock allocationCounterStartVal tracingSpanRef >> andThen onFinish,
-        trackAnalyticsEventIO = trackEventIO
+        trackAnalyticsEventIO = trackEventIO,
+        sessionIdRef = sidRef
       }
 
 -- | Helper that creates a handler that does nothing. This is intended to power
@@ -643,6 +666,7 @@ nullHandler = do
     { requestId = "",
       startChildTracingSpan = \_ -> pure nullHandler,
       startNewRoot = \_ -> pure nullHandler,
+      sessionIdRef = nullSessionIdRef,
       setTracingSpanDetailsIO = \_ -> pure (),
       setTracingSpanSummaryIO = \_ -> pure (),
       markTracingSpanFailedIO = pure (),
@@ -871,8 +895,11 @@ rootTracingSpanIO ::
   IO a
 rootTracingSpanIO requestId trackEventIO onFinish name runIO = do
   clock' <- mkClock
+  -- Each request gets its own session-id ref. Children share it via
+  -- closure capture inside `mkHandler`.
+  sidRef <- IORef.newIORef Nothing
   Exception.bracketWithError
-    (Stack.withFrozenCallStack mkHandler requestId clock' trackEventIO onFinish Nothing name)
+    (Stack.withFrozenCallStack mkHandler requestId clock' trackEventIO sidRef onFinish Nothing name)
     (Prelude.flip finishTracingSpan)
     runIO
 

nri-prelude/tests/LogSpec.hs

@@ -190,12 +190,14 @@ instance Exception.Exception TestException
 newHandler :: (Stack.HasCallStack) => Prelude.IO (Prelude.IO [Internal.TracingSpan], Internal.LogHandler)
 newHandler = do
   recordedTracingSpans <- IORef.newIORef []
+  sidRef <- IORef.newIORef Nothing
   handler <-
     Stack.withFrozenCallStack
       Internal.mkHandler
       ""
       (Internal.Clock (Prelude.pure 0))
       Internal.silentTrack
+      sidRef
       (\span -> IORef.modifyIORef recordedTracingSpans (\cs -> cs ++ Internal.children span))
       Nothing
       ""

nri-prelude/tests/PlatformSpec.hs

@@ -60,7 +60,41 @@ tests =
           Ok () -> Expect.pass
           Err _ -> Expect.fail "trackEvent task failed"
         observed <- Expect.fromIO (IORef.readIORef ref)
-        observed |> Expect.equal [event]
+        observed |> Expect.equal [event],
+      test "trackEvent stamps the current request's session_id onto the payload" <| \_ -> do
+        ref <- Expect.fromIO (IORef.newIORef [])
+        let track v = IORef.atomicModifyIORef' ref (\xs -> (v : xs, ()))
+        let event = Aeson.object ["kind" Aeson..= ("LessonStarted" :: Text)]
+        let expected = Aeson.object ["kind" Aeson..= ("LessonStarted" :: Text), "session_id" Aeson..= ("sess-42" :: Text)]
+        Expect.fromIO
+          <| Platform.rootTracingSpanIO "test-req" track (\_ -> Prelude.pure ()) "root"
+          <| \log -> do
+            Platform.setSessionIdIO log (Just "sess-42")
+            _ <- Task.attempt log (Platform.Analytics.Internal.trackEvent event)
+            Prelude.pure ()
+        observed <- Expect.fromIO (IORef.readIORef ref)
+        observed |> Expect.equal [expected],
+      test "Platform.sessionId reflects the value written by setSessionIdIO" <| \_ -> do
+        result <-
+          Expect.fromIO
+            <| Platform.rootTracingSpanIO "test-req" Platform.silentTrack (\_ -> Prelude.pure ()) "root"
+            <| \log -> do
+              Platform.setSessionIdIO log (Just "sess-7")
+              Task.attempt log Platform.sessionId
+        case result of
+          Ok mSid -> mSid |> Expect.equal (Just "sess-7")
+          Err _ -> Expect.fail "sessionId task failed",
+      test "child handlers see the session id set on the root" <| \_ -> do
+        result <-
+          Expect.fromIO
+            <| Platform.rootTracingSpanIO "test-req" Platform.silentTrack (\_ -> Prelude.pure ()) "root"
+            <| \log -> do
+              Platform.setSessionIdIO log (Just "sess-shared")
+              child <- Platform.Internal.startChildTracingSpan log "child-span"
+              Task.attempt child Platform.sessionId
+        case result of
+          Ok mSid -> mSid |> Expect.equal (Just "sess-shared")
+          Err _ -> Expect.fail "sessionId task failed"
     ]
 
 newtype CustomTracingSpanDetails = CustomTracingSpanDetails Text