Trace.hs

{-# LANGUAGE ScopedTypeVariables #-}

-- | The 'withTrace' wrapper and supporting helpers.
--
-- After EP-3, the trace bridge is stream-shaped at the core:
-- 'withTraceStream' returns a 'Stream IO AssistantMessageEvent'
-- that side-effects 'CallStarted' / 'CallFinished' / 'CallFailed'
-- events to a user-supplied 'TraceSink' as the stream's lifecycle
-- unfolds. 'withTrace' is the synchronous draining wrapper:
-- @withTrace sink m ctx opts = Stream.fold (reassembleResponse m)
-- (withTraceStream sink m ctx opts)@.
--
-- Per-call plumbing: open a 'Chan' of @Maybe TraceEvent@, fork a
-- worker thread that drains the channel through the sink's fold,
-- push 'CallStarted' eagerly (before the first
-- 'AssistantMessageEvent' is emitted), then watch for the stream's
-- terminal event ('EventDone' or 'EventError') and push the
-- matching 'CallFinished' / 'CallFailed' before yielding the
-- terminal event to the consumer. Cleanup ('Nothing' sentinel on
-- the channel + 'takeMVar' on the worker) is idempotent and runs
-- through 'Stream.finallyIO' so an early-aborting consumer eventually
-- records a synthetic 'CallFailed' and never leaks the worker. Sink
-- exceptions are captured by the worker and reported once on stderr
-- during cleanup; they do not propagate into the provider call.
module Baikai.Trace
  ( -- * Re-exports
    TraceEvent (..),
    TraceSink (..),

    -- * Wrappers
    withTrace,
    withTraceWith,
    withTraceStream,
    withTraceStreamWith,
    runRequestWith,
    runRequestWithRegistry,

    -- * Helpers
    newEventId,
    summarizeContext,
  )
where

import Baikai.Context (Context)
import Baikai.Cost (usdAsScientific)
import Baikai.Cost qualified as Cost
import Baikai.Cost.Log
  ( CallLogEntry (..),
    CallLogHandle,
    appendEntry,
    summarizeContext,
  )
import Baikai.Message (AssistantPayload (..), Message (..))
import Baikai.Model (Model)
import Baikai.Options (Options)
import Baikai.Prelude
import Baikai.Provider.Registry (ProviderRegistry, globalProviderRegistry)
import Baikai.Response (Response)
import Baikai.Stream (reassembleResponse, streamRequestWith)
import Baikai.Stream.Event (AssistantMessageEvent (..), TerminalPayload (..))
import Baikai.Trace.Event (TraceEvent (..))
import Baikai.Trace.Sink (TraceSink (..))
import Baikai.Usage (Usage)
import Baikai.Usage qualified as Usage
import Control.Concurrent (forkIO)
import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, takeMVar)
import Control.Exception (SomeException, displayException, try)
import Control.Monad (forM_, unless)
import Control.Monad.IO.Unlift (MonadUnliftIO, withRunInIO)
import Data.Bits (unsafeShiftL, (.&.), (.|.))
import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef, writeIORef)
import Data.Maybe (fromMaybe)
import Data.Text qualified as Text
import Data.Time (UTCTime, diffUTCTime, getCurrentTime)
import Data.Time.Clock.POSIX (getPOSIXTime)
import Data.Word (Word64)
import Foreign.StablePtr (StablePtr, freeStablePtr, newStablePtr)
import Numeric (showHex)
import Streamly.Data.Stream (Stream)
import Streamly.Data.Stream qualified as Stream
import System.IO (hPutStrLn, stderr)
import System.IO.Unsafe (unsafePerformIO)

-- ============================================================
-- Stream-shaped trace bridge
-- ============================================================

-- | Decorate a streaming provider call with structured tracing
-- events emitted to the supplied 'TraceSink'.
--
-- The returned stream yields the same 'AssistantMessageEvent's as
-- 'Baikai.Stream.streamRequest'. As a side effect: one 'CallStarted'
-- event is pushed to the sink before the first
-- 'AssistantMessageEvent' is observed, and one 'CallFinished' (on
-- 'EventDone') or 'CallFailed' (on 'EventError') is pushed before
-- the terminal event is yielded to the consumer. Intermediate delta
-- events are not traced — emitting one trace per token would
-- explode trace volume.
withTraceStream ::
  TraceSink ->
  Model ->
  Context ->
  Options ->
  Stream IO AssistantMessageEvent
withTraceStream = withTraceStreamWith globalProviderRegistry

-- | Decorate a streaming provider call through an explicit registry handle.
withTraceStreamWith ::
  ProviderRegistry ->
  TraceSink ->
  Model ->
  Context ->
  Options ->
  Stream IO AssistantMessageEvent
withTraceStreamWith reg (TraceSink sinkFold) m ctx opts =
  Stream.concatEffect $ do
    state <- newTraceState
    let c = state ^. #chan
        d = state ^. #done
    _ <-
      forkIO $ do
        let stepDrain () = do
              msg <- readChan c
              pure (fmap (\e -> (e, ())) msg)
        result <-
          try
            ( Stream.unfoldrM stepDrain ()
                & Stream.fold sinkFold
            ) ::
            IO (Either SomeException ())
        case result of
          Left e -> writeIORef (state ^. #sinkError) (Just e)
          Right () -> pure ()
        putMVar d ()
    eid <- newEventId
    start <- getCurrentTime
    writeChan c $
      Just
        CallStarted
          { eventId = eid,
            timestamp = start,
            provider = m ^. #provider,
            model = m ^. #modelId,
            maxTokens = resolvedMaxTokens m opts,
            promptSummary = summarizeContext ctx
          }
    pure $
      Stream.finallyIO
        (finalizeTrace state eid start m)
        (Stream.mapM (traceEvent state eid start m) (streamRequestWith reg m ctx opts))

-- | Synchronous trace wrapper. Drains 'withTraceStream' into a
-- 'Response' through 'reassembleResponse'.
--
-- Unlike the EP-2 'withTrace' (which re-threw the producer's
-- exception), this implementation never throws for producer-side
-- failures: errors flow through the stream as a terminal
-- 'EventError' and the drained 'Response' carries
-- @stopReason = ErrorReason@ plus 'errorMessage'. The masterplan's
-- Vision & Scope section commits to "partial output is always
-- recoverable" and the plan's Decision Log records that producer
-- failures must surface as response data, not exceptions.
-- Downstream-of-the-fold exceptions (e.g. an 'appendEntry' that
-- fails) still propagate unchanged.
withTrace ::
  (MonadUnliftIO m) =>
  TraceSink -> Model -> Context -> Options -> m Response
withTrace = withTraceWith globalProviderRegistry

-- | Synchronous trace wrapper that dispatches through an explicit provider
-- registry handle.
withTraceWith ::
  (MonadUnliftIO m) =>
  ProviderRegistry ->
  TraceSink ->
  Model ->
  Context ->
  Options ->
  m Response
withTraceWith reg sink model ctx opts =
  withRunInIO $ \_ ->
    Stream.fold
      (reassembleResponse model)
      (withTraceStreamWith reg sink model ctx opts)

-- ============================================================
-- Per-call trace state
-- ============================================================

data TraceState = TraceState
  { chan :: !(Chan (Maybe TraceEvent)),
    done :: !(MVar ()),
    closed :: !(IORef Bool),
    sinkError :: !(IORef (Maybe SomeException)),
    terminalSent :: !(IORef Bool),
    stableRoot :: !(IORef (Maybe (StablePtr TraceState)))
  }
  deriving stock (Generic)

newTraceState :: IO TraceState
newTraceState = do
  c <- newChan
  d <- newEmptyMVar
  r <- newIORef False
  e <- newIORef Nothing
  t <- newIORef False
  root <- newIORef Nothing
  let state = TraceState {chan = c, done = d, closed = r, sinkError = e, terminalSent = t, stableRoot = root}
  sp <- newStablePtr state
  writeIORef root (Just sp)
  pure state

finalizeTrace :: TraceState -> Text -> UTCTime -> Model -> IO ()
finalizeTrace s eid start m = do
  alreadyClosed <-
    atomicModifyIORef' (s ^. #closed) (\b -> (True, b))
  unless alreadyClosed $ do
    sent <- readIORef (s ^. #terminalSent)
    unless sent $ do
      now <- getCurrentTime
      writeChan (s ^. #chan) $
        Just
          CallFailed
            { eventId = eid,
              timestamp = now,
              provider = m ^. #provider,
              model = m ^. #modelId,
              latencyMs = millisBetween start now,
              errorMessage = "aborted: stream consumer stopped before the terminal event"
            }
    writeChan (s ^. #chan) Nothing
    takeMVar (s ^. #done)
    reportSinkError s
    releaseStableRoot s

releaseStableRoot :: TraceState -> IO ()
releaseStableRoot s = do
  msp <- atomicModifyIORef' (s ^. #stableRoot) (\sp -> (Nothing, sp))
  forM_ msp freeStablePtr

reportSinkError :: TraceState -> IO ()
reportSinkError s = do
  merr <- readIORef (s ^. #sinkError)
  forM_ merr $ \e ->
    hPutStrLn
      stderr
      ("baikai: trace sink failed; trace events for this call were dropped: " <> displayException e)

traceEvent ::
  TraceState ->
  Text ->
  UTCTime ->
  Model ->
  AssistantMessageEvent ->
  IO AssistantMessageEvent
traceEvent state eid start m ev = do
  case ev of
    EventDone TerminalPayload {message = msg} -> do
      now <- getCurrentTime
      let latency = millisBetween start now
          mu = assistantUsageFromMsg msg
          meaningfulCost = maybe False (\u -> usdRat (Usage.cost u) > 0) mu
          finished =
            CallFinished
              { eventId = eid,
                timestamp = now,
                provider = m ^. #provider,
                model = m ^. #modelId,
                latencyMs = latency,
                inputTokens = fmap Usage.inputTokens mu,
                outputTokens = fmap Usage.outputTokens mu,
                usd =
                  if meaningfulCost
                    then fmap (usdAsScientific . Usage.cost) mu
                    else Nothing
              }
      writeChan (state ^. #chan) (Just finished)
      writeIORef (state ^. #terminalSent) True
      finalizeTrace state eid start m
    EventError TerminalPayload {message = msg} -> do
      now <- getCurrentTime
      let latency = millisBetween start now
          errMsg = case msg of
            AssistantMessage AssistantPayload {errorMessage = Just t} -> t
            _ -> "stream terminated with EventError"
          failed =
            CallFailed
              { eventId = eid,
                timestamp = now,
                provider = m ^. #provider,
                model = m ^. #modelId,
                latencyMs = latency,
                errorMessage = errMsg
              }
      writeChan (state ^. #chan) (Just failed)
      writeIORef (state ^. #terminalSent) True
      finalizeTrace state eid start m
    _ -> pure ()
  pure ev

-- ============================================================
-- Cost-log convenience wrapper
-- ============================================================

-- | Combine tracing with call-log persistence in one call. Traces
-- the streaming call, drains it into a 'Response', then appends a
-- 'CallLogEntry' to the given handle.
runRequestWith ::
  (MonadUnliftIO m) =>
  TraceSink ->
  CallLogHandle ->
  Model ->
  Context ->
  Options ->
  m Response
runRequestWith = runRequestWithRegistry globalProviderRegistry

-- | Combine tracing with call-log persistence while dispatching through an
-- explicit provider registry handle.
runRequestWithRegistry ::
  (MonadUnliftIO m) =>
  ProviderRegistry ->
  TraceSink ->
  CallLogHandle ->
  Model ->
  Context ->
  Options ->
  m Response
runRequestWithRegistry reg sink h m ctx opts = do
  resp <- withTraceWith reg sink m ctx opts
  now <- liftIO getCurrentTime
  let mu = assistantUsage resp
      meaningfulCost = maybe False (\u -> usdRat (Usage.cost u) > 0) mu
      entry =
        CallLogEntry
          { timestamp = now,
            provider = m ^. #provider,
            model = m ^. #modelId,
            inputTokens = mu >>= positiveNat . Usage.inputTokens,
            outputTokens = mu >>= positiveNat . Usage.outputTokens,
            cachedInputTokens = mu >>= positiveNat . Usage.cacheReadTokens,
            reasoningTokens = mu >>= Usage.reasoningTokens,
            usd =
              if meaningfulCost
                then fmap (usdAsScientific . Usage.cost) mu
                else Nothing,
            latencyMs = resp ^. #latencyMs,
            promptSummary = summarizeContext ctx
          }
  appendEntry h entry
  pure resp

-- ============================================================
-- Internal helpers
-- ============================================================

-- | Project the assistant turn's 'Usage' out of a response.
assistantUsage :: Response -> Maybe Usage
assistantUsage resp = Just ((resp ^. #message) ^. #usage)

assistantUsageFromMsg :: Message -> Maybe Usage
assistantUsageFromMsg = \case
  AssistantMessage AssistantPayload {usage = u} -> Just u
  _ -> Nothing

usdRat :: Cost.Cost -> Rational
usdRat = Cost.usd

positiveNat :: Natural -> Maybe Natural
positiveNat 0 = Nothing
positiveNat n = Just n

resolvedMaxTokens :: Model -> Options -> Natural
resolvedMaxTokens m opts = fromMaybe (m ^. #maxOutputTokens) (opts ^. #maxTokens)

millisBetween :: UTCTime -> UTCTime -> Int
millisBetween a b = round (realToFrac (diffUTCTime b a) * (1000 :: Double))

-- ============================================================
-- Event id
-- ============================================================

-- | Generate a 16-character lowercase hexadecimal event id. The high
-- 32 bits are derived from process-start POSIX seconds and the low
-- 32 bits are a process-local counter, so ids are unique within a
-- process for 2^32 calls.
newEventId :: IO Text
newEventId = do
  n <- atomicModifyIORef' eventCounter (\k -> (k + 1, k))
  let raw :: Word64
      raw =
        (fromIntegral eventBase .&. 0xFFFFFFFF) `unsafeShiftL` 32
          .|. (fromIntegral n .&. 0xFFFFFFFF)
      hex = showHex raw ""
      padded = replicate (16 - length hex) '0' <> hex
  pure (Text.pack padded)

eventCounter :: IORef Word
eventCounter = unsafePerformIO (newIORef 0)
{-# NOINLINE eventCounter #-}

eventBase :: Word
eventBase = unsafePerformIO $ do
  t <- getPOSIXTime
  pure (fromIntegral (floor t :: Integer))
{-# NOINLINE eventBase #-}