From 3eba39c4ae1021d4330d7924c0d6362a206d4395 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arthur=20Bernardi=20Jord=C3=A3o?= Date: Fri, 8 May 2026 14:38:04 -0300 Subject: [PATCH] nri-prelude: thread session id through LogHandler Adds a per-request session-id slot to LogHandler with the same plumbing pattern as trackAnalyticsEventIO: an IORef created by rootTracingSpanIO and shared across child / new-root handlers via closure capture in mkHandler. Public surface: - Platform.sessionId :: Task e (Maybe Text) - Platform.setSessionIdIO :: LogHandler -> Maybe Text -> IO () The intended use is a WAI middleware on the application root: read an incoming X-NRI-Session-Id header, call setSessionIdIO once on the per-request LogHandler, and let every Analytics.track call inside the request stamp it automatically. Platform.Analytics.Internal.trackEvent reads the ref and shallow-merges { "session_id": ... } onto the event's JSON before invoking the analytics callback, so per-event schemas don't need a session_id field. Mirrors how event_id/event_timestamp are stamped centrally in nri-analytics. Tests cover four cases: - trackEvent stamps session_id when set - existing trackEvent test (no session id) still passes unchanged - sessionId reflects setSessionIdIO writes - child handlers see the session id set on the root Co-Authored-By: Claude Opus 4.7 (1M context) --- nri-prelude/src/Platform.hs | 23 +++++++++++ .../src/Platform/Analytics/Internal.hs | 36 ++++++++++++----- nri-prelude/src/Platform/Internal.hs | 39 ++++++++++++++++--- nri-prelude/tests/LogSpec.hs | 2 + nri-prelude/tests/PlatformSpec.hs | 36 ++++++++++++++++- 5 files changed, 119 insertions(+), 17 deletions(-) diff --git a/nri-prelude/src/Platform.hs b/nri-prelude/src/Platform.hs index d3caf61f..19070c28 100644 --- a/nri-prelude/src/Platform.hs +++ b/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 diff --git a/nri-prelude/src/Platform/Analytics/Internal.hs b/nri-prelude/src/Platform/Analytics/Internal.hs index 46b32a28..0d34a7e7 100644 --- a/nri-prelude/src/Platform/Analytics/Internal.hs +++ b/nri-prelude/src/Platform/Analytics/Internal.hs @@ -14,6 +14,7 @@ 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 @@ -21,18 +22,33 @@ 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 diff --git a/nri-prelude/src/Platform/Internal.hs b/nri-prelude/src/Platform/Internal.hs index 213bd5dc..df9eda19 100644 --- a/nri-prelude/src/Platform/Internal.hs +++ b/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,7 +580,15 @@ 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` @@ -587,6 +596,16 @@ data LogHandler = LogHandler 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 diff --git a/nri-prelude/tests/LogSpec.hs b/nri-prelude/tests/LogSpec.hs index d75d4c8f..c394ba27 100644 --- a/nri-prelude/tests/LogSpec.hs +++ b/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 "" diff --git a/nri-prelude/tests/PlatformSpec.hs b/nri-prelude/tests/PlatformSpec.hs index 2908dffb..4fabfb23 100644 --- a/nri-prelude/tests/PlatformSpec.hs +++ b/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