SSE.hs

{-# LANGUAGE OverloadedStrings #-}

-- |
-- Module      : Database.Bloodhound.Internal.Utils.SSE
-- Description : Server-Sent Events parser over an http-client BodyReader
--
-- A small, dependency-light Server-Sent Events (SSE) parser that pulls
-- bytes incrementally from a 'Network.HTTP.Client.BodyReader' (the
-- streaming response body exposed by 'Database.Bloodhound.Client.Cluster.withStreamingResponse').
-- It is /not/ a general-purpose SSE server; it is the minimum needed
-- to consume SSE-style responses from the OpenSearch ML Commons
-- Execute Stream Agent API
-- (@POST /_plugins/_ml/agents/{agent_id}/_execute/stream@), whose
-- wire format is one @data: <JSON>\\n\\n@ frame per agent chunk.
--
-- The parser implements the subset of the SSE specification that
-- matters for parsing a byte stream into events:
--
--   * Lines are terminated by @\\n@, optionally with a preceding @\\r@.
--   * An event is terminated by a blank line (the canonical @\\n\\n@
--     separator, tolerating @\\r\\n\\r\\n@).
--   * Lines starting with @:@ are comments and are ignored.
--   * A line of the form @field: value@ sets the named field; the
--     leading single space after the colon is stripped (per spec),
--     further spaces are preserved. A line with no colon sets the
--     field to the empty string with an empty value.
--   * Multiple @data:@ lines within a single event are joined with
--     @\\n@ to form the final 'sseData' (per spec).
--   * The @event@, @id@, and @retry@ fields are parsed (so a future
--     endpoint that emits them round-trips), but the Execute Stream
--     Agent API only ever emits @data:@.
--
-- The parser yields 'SSEEvent' values one at a time; the caller
-- decides when to stop (e.g. on a chunk whose JSON carries an
-- @is_last@ flag, or an AG-UI @RUN_FINISHED@ discriminator). There is
-- no synthetic end-of-stream event: when the 'BodyReader' is
-- exhausted, 'readSSEEvent' returns 'Nothing'.
module Database.Bloodhound.Internal.Utils.SSE
  ( SSEEvent (..),
    SSEState,
    newSSEState,
    readSSEEvent,
    sseSource,
  )
where

import Conduit (ConduitT, MonadResource, liftIO, yield)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef)
import Data.Maybe (isJust)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Data.Text.Encoding.Error (lenientDecode)
import Network.HTTP.Client (BodyReader, brRead)
import Numeric (readDec)
import Prelude hiding (head, tail)

-- | A single parsed Server-Sent Event. Mirrors the fields the spec
-- defines; 'sseEvent' is 'Nothing' when no @event:@ line was present
-- (the default-event case, which is what the Execute Stream Agent API
-- always produces).
data SSEEvent = SSEEvent
  { sseEvent :: Maybe Text,
    sseData :: Text,
    sseId :: Maybe Text,
    sseRetry :: Maybe Int
  }
  deriving stock (Eq, Show)

-- | The empty event used to begin accumulating a new frame.
emptyEvent :: SSEEvent
emptyEvent = SSEEvent Nothing "" Nothing Nothing

-- | Whether an accumulated event carries any content (at least one
-- field was set). A blank-line terminator on an empty event is a
-- keep-alive heartbeat; the parser still dispatches such events so
-- the caller can observe heartbeats if it cares, but at end-of-stream
-- a trailing empty event is dropped.
eventHasContent :: SSEEvent -> Bool
eventHasContent ev =
  not (T.null (sseData ev))
    || isJust (sseEvent ev)
    || isJust (sseId ev)
    || isJust (sseRetry ev)

-- | Mutable parser state: the unconsumed byte buffer (leftover bytes
-- after the last dispatched event) paired with the underlying
-- 'BodyReader'. Carrying the buffer in an 'IORef' lets 'readSSEEvent'
-- retain leftover bytes across calls without re-reading from the
-- network.
data SSEState = SSEState
  { sseStateBuffer :: IORef ByteString,
    sseStateReader :: BodyReader
  }

-- | Allocate a fresh parser state over a 'BodyReader'.
newSSEState :: BodyReader -> IO SSEState
newSSEState reader = SSEState <$> newIORef BS.empty <*> pure reader

-- | Pull the next complete SSE event from the stream. Returns
-- 'Nothing' when the body is fully consumed and no partial event with
-- content remains. See module docs for the parsing rules.
readSSEEvent :: SSEState -> IO (Maybe SSEEvent)
readSSEEvent st = go emptyEvent
  where
    go event = do
      buf <- readIORef (sseStateBuffer st)
      case breakLine buf of
        Just (line, rest) -> do
          writeIORef (sseStateBuffer st) rest
          case applyLine line event of
            Dispatch ev -> pure (Just ev)
            Continue ev' -> go ev'
        Nothing -> do
          chunk <- brRead (sseStateReader st)
          if BS.null chunk
            then do
              -- End of stream. The buffer may still hold a trailing
              -- partial line (one with no terminating @\\n@); feed it
              -- through 'applyLine' as a final line before deciding
              -- whether to dispatch, so an unterminated @data: ...@
              -- frame is not lost.
              leftover <- readIORef (sseStateBuffer st)
              let finalEvent =
                    if BS.null leftover
                      then event
                      else case applyLine leftover event of
                        Dispatch ev -> ev
                        Continue ev' -> ev'
              if eventHasContent finalEvent
                then pure (Just finalEvent)
                else pure Nothing
            else do
              modifyIORef' (sseStateBuffer st) (<> chunk)
              go event

-- | The result of interpreting a single line: either the accumulated
-- event is complete and should be dispatched, or the accumulator
-- should be updated and parsing continue.
data LineResult
  = Dispatch SSEEvent
  | Continue SSEEvent

-- | Interpret one SSE line against the current accumulator. A blank
-- line dispatches (per spec); a comment line (@:@) is a no-op
-- continue; any other line updates the named field and continues.
applyLine :: ByteString -> SSEEvent -> LineResult
applyLine line event
  | BS.null line = Dispatch event
  | BS.head line == colon = Continue event -- comment
  | otherwise = Continue (applyField field value event)
  where
    (field, value) = splitField line
    colon = toEnum 58

-- | Split a line into @(field, value)@ on the first colon. A leading
-- single space in the value is stripped per spec. A line with no colon
-- has the field as the whole line and an empty value.
splitField :: ByteString -> (ByteString, ByteString)
splitField line =
  case BS.elemIndex colon line of
    Just i ->
      let rawValue = BS.drop (i + 1) line
          value = stripOneSpace rawValue
       in (BS.take i line, value)
    Nothing -> (line, BS.empty)
  where
    colon = toEnum 58

stripOneSpace :: ByteString -> ByteString
stripOneSpace s = case BS.uncons s of
  Just (c, rest) | c == toEnum 32 -> rest -- ' '
  _ -> s

applyField :: ByteString -> ByteString -> SSEEvent -> SSEEvent
applyField field value event
  | field == "data" =
      event
        { sseData =
            if T.null (sseData event)
              then decodeUtf8Lenient value
              else sseData event <> "\n" <> decodeUtf8Lenient value
        }
  | field == "event" = event {sseEvent = Just (decodeUtf8Lenient value)}
  | field == "id" = event {sseId = Just (decodeUtf8Lenient value)}
  | field == "retry" = event {sseRetry = parseRetry value}
  | otherwise = event

decodeUtf8Lenient :: ByteString -> Text
decodeUtf8Lenient = TE.decodeUtf8With lenientDecode

parseRetry :: ByteString -> Maybe Int
parseRetry v = case readDec (charStr v) of
  (n, _) : _ -> Just n
  [] -> Nothing
  where
    -- readDec works on String; convert losslessly.
    charStr = map (toEnum . fromIntegral) . BS.unpack

-- | Split the buffer on the first line terminator (@\\n@, optionally
-- preceded by @\\r@). Returns the line (without the terminator) and
-- the remaining buffer, or 'Nothing' if no terminator is present yet.
breakLine :: ByteString -> Maybe (ByteString, ByteString)
breakLine buf =
  case BS.elemIndex newline buf of
    Nothing -> Nothing
    Just i ->
      let termStart = if i > 0 && BS.index buf (i - 1) == carriage then i - 1 else i
          line = BS.take termStart buf
          rest = BS.drop (i + 1) buf
       in Just (line, rest)
  where
    newline = toEnum 10
    carriage = toEnum 13

-- | A conduit source over a 'BodyReader' that yields 'SSEEvent's
-- until the body is exhausted. The caller (typically wrapped in
-- 'Data.Conduit.runConduitRes') decides when to stop consuming based
-- on event content; the source itself terminates when the body ends.
sseSource :: (MonadResource m) => BodyReader -> ConduitT () SSEEvent m ()
sseSource reader = do
  st <- liftIO (newSSEState reader)
  loop st
  where
    loop st = do
      next <- liftIO (readSSEEvent st)
      case next of
        Just ev -> yield ev >> loop st
        Nothing -> pure ()

-- $setup
-- (No doctest setup; the module is exercised via @Test.MLModelSpec@.)