StreamingSpec.hs

{-# OPTIONS_GHC -Wno-missing-fields #-}

module LLM.StreamingSpec (spec) where

import Control.Concurrent (MVar, forkIO, newEmptyMVar, putMVar, takeMVar, threadDelay)
import Control.Exception (finally, throwIO, try)
import Control.Monad (when)
import Control.Retry (RetryPolicyM (getRetryPolicyM), RetryStatus (rsIterNumber), defaultRetryStatus)
import Data.Aeson (object, (.=))
import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef, writeIORef)
import Data.Text (Text)
import GHC.Clock (getMonotonicTime)
import LLM.Core.Abort (AbortSignal, abort, newAbortSignal)
import LLM.Core.LLMProvider (LLMProvider (..), toGateway)
import LLM.Providers.OpenAI (openAIProvider)
import Network.HTTP.Client qualified as HC
import Network.HTTP.Req (HttpException (..))
import LLM.Core.Types
  ( ChatResponse (..),
    LLMError (..),
    LLMGateway (..),
    LLMTextResult,
    supportedImageMediaTypes,
    StreamEvent (..),
    Turn (..),
    assistantTurn,
    mkToolCall,
    pattern UserTurn,
    textPart,
    toolCallPart,
  )
import LLM.Core.Usage (PricingInfo (..), defaultPricingInfo, mkUsage)
import LLM.Generate.Generate (generateTextLLM, generateTextWithFallbacks, streamTextLLM, streamTextWithFallbacks)
import LLM.Generate.GenerateUtils (llmHooks)
import LLM.Generate.Logger (noHooks)
import LLM.Generate.ModelConfig (ModelConfig (..), ModelWithFallbacks (..), defaultModelCapabilities, modelRetryPolicy)
import LLM.Generate.Types
  ( GenRequest (..),
    GenerateError (..),
    RoundTextRole (..),
    StreamChunk (..),
  )
import System.Timeout (timeout)
import Test.Hspec (Spec, describe, it, shouldBe, shouldReturn)

spec :: Spec
spec = describe "Streaming" $ do
  describe "streamTextLLM phase routing" $ do
    it "classifies plain text as answer on finalize" $ do
      let gw =
            streamGateway
              [StreamDelta "hello"]
              (ChatResponse "hello" [textPart "hello"] (Just (mkUsage 1 1)) Nothing)
      chunks <- runStream gw []
      reverse chunks
        `shouldBe` [ TextDelta "hello",
                     RoundTextRoleCommitted AnswerRole
                   ]

    it "routes text to answer after a tool turn" $ do
      let prior = [UserTurn "q", assistantTurn "" Nothing [mkToolCall "1" "t" (object [])], ToolTurn []]
          gw =
            streamGateway
              [StreamDelta "done"]
              (ChatResponse "done" [textPart "done"] (Just (mkUsage 1 1)) Nothing)
      chunks <- runStream gw prior
      reverse chunks `shouldBe` [AnswerDelta "done"]

    it "routes preamble text and tool calls before answer" $ do
      let tc = mkToolCall "c1" "grep" (object ["pattern" .= ("x" :: Text)])
          gw =
            streamGateway
              [StreamDelta "searching", StreamToolCall tc]
              (ChatResponse "" [toolCallPart tc] (Just (mkUsage 2 0)) Nothing)
      chunks <- runStream gw [UserTurn "find x"]
      reverse chunks
        `shouldBe` [ TextDelta "searching",
                     RoundTextRoleCommitted PreambleRole,
                     StreamToolCallChunk tc
                   ]

    it "forwards reasoning deltas unchanged" $ do
      let gw =
            streamGateway
              [StreamReasoningDelta "think"]
              (ChatResponse "ok" [textPart "ok"] Nothing Nothing)
      chunks <- runStream gw []
      reverse chunks `shouldBe` [ReasoningDelta "think", RoundTextRoleCommitted AnswerRole]

  describe "request timeout" $ do
    it "applies the timeout per attempt, so a retry can succeed" $ do
      attempts <- newIORef (0 :: Int)
      let resp = ChatResponse "ok" [textPart "ok"] Nothing Nothing
          gw =
            (streamGateway [] resp)
              { gwGenerateText = \_ _ -> do
                  n <- atomicModifyIORef' attempts (\k -> (k + 1, k))
                  when (n == 0) $ threadDelay 1_000_000
                  pure (Right resp)
              }
          mc = (testModelConfig gw) {mcRequestTimeout = Just 100, mcRetryCount = 1}
      r <- generateTextLLM (testGenRequest []) mc
      fmap (.respText) r `shouldBe` Right "ok"
      readIORef attempts >>= (`shouldBe` 2)

    it "streaming: does not cut off a stream that started within the timeout" $ do
      let resp = ChatResponse "abc" [textPart "abc"] Nothing Nothing
          gw =
            (streamGateway [] resp)
              { gwStreamText = \_ _ onEvent -> do
                  onEvent (StreamDelta "a")
                  threadDelay 150_000
                  onEvent (StreamDelta "b")
                  threadDelay 150_000
                  onEvent (StreamDelta "c")
                  pure (Right resp)
              }
          mc = (testModelConfig gw) {mcRequestTimeout = Just 100}
      ref <- newIORef []
      r <- streamTextLLM (\c -> modifyIORef' ref (c :)) (testGenRequest []) mc
      fmap (.respText) r `shouldBe` Right "abc"
      chunks <- reverse <$> readIORef ref
      take 3 chunks `shouldBe` [TextDelta "a", TextDelta "b", TextDelta "c"]

    it "streaming: times out when no event arrives in time" $ do
      let resp = ChatResponse "late" [textPart "late"] Nothing Nothing
          gw =
            (streamGateway [] resp)
              { gwStreamText = \_ _ onEvent -> do
                  threadDelay 1_000_000
                  onEvent (StreamDelta "late")
                  pure (Right resp)
              }
          mc = (testModelConfig gw) {mcRequestTimeout = Just 100}
      r <- streamTextLLM (\_ -> pure ()) (testGenRequest []) mc
      fmap (.respText) r `shouldBe` Left TimeoutError

  describe "stream callback exceptions" $ do
    it "rethrows the callback's own HttpException after one attempt" $ do
      attempts <- newIORef (0 :: Int)
      let provider =
            (openAIProvider "key")
              { openStream = \_ k -> do
                  modifyIORef' attempts (+ 1)
                  chunks <- newIORef ["data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n", "data: [DONE]\n\n"]
                  k (atomicModifyIORef' chunks (\case [] -> ([], ""); c : cs -> (cs, c)))
              }
          mc = (testModelConfig (toGateway provider)) {mcRetryCount = 3}
          callerFailure = VanillaHttpException (HC.HttpExceptionRequest HC.defaultRequest HC.ConnectionClosed)
          -- Throw only for text, not for the 'StreamAttemptReset' a retry
          -- would emit first, so a retry would show as a second attempt.
          onChunk = \case
            TextDelta _ -> throwIO callerFailure
            _ -> pure ()
      r <- try (streamTextLLM onChunk (testGenRequest [UserTurn "q"]) mc)
      case r of
        Left (VanillaHttpException (HC.HttpExceptionRequest _ HC.ConnectionClosed)) -> pure ()
        Left e -> fail ("unexpected exception: " <> show e)
        Right res -> fail ("expected the exception, got " <> show res)
      readIORef attempts `shouldReturn` 1

    it "still ends a provider stream with Aborted when the signal fires in the callback" $ do
      sig <- newAbortSignal
      let provider =
            (openAIProvider "key")
              { openStream = \_ k -> do
                  chunks <-
                    newIORef
                      [ "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n",
                        "data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n",
                        "data: [DONE]\n\n"
                      ]
                  k (atomicModifyIORef' chunks (\case [] -> ([], ""); c : cs -> (cs, c)))
              }
      r <- streamTextLLM (\_ -> abort sig) (testGenRequest [UserTurn "q"]) {grAbortSignal = Just sig} (testModelConfig (toGateway provider))
      fmap (.respText) r `shouldBe` Left Aborted

  describe "abort signal" $ do
    it "does not call the model when the signal fired before the call" $ do
      (gw, calls) <- countingGateway (Right (ChatResponse "hi" [textPart "hi"] Nothing Nothing))
      sig <- newAbortSignal
      abort sig
      r <- generateTextWithFallbacks (testGenRequest []) {grAbortSignal = Just sig} (ModelWithFallbacks (testModelConfig gw) [])
      fmap (.respText) r `shouldBe` Left GErrAborted
      readIORef calls `shouldReturn` 0

    it "stops retrying once the signal fired" $ do
      sig <- newAbortSignal
      (gw, calls) <- countingGateway (Left (HttpError 503 "overloaded"))
      let gw' = gw {gwGenerateText = \h req -> abort sig >> gw.gwGenerateText h req}
          mc = (testModelConfig gw') {mcRetryCount = 3}
      r <- generateTextWithFallbacks (testGenRequest []) {grAbortSignal = Just sig} (ModelWithFallbacks mc [])
      fmap (.respText) r `shouldBe` Left GErrAborted
      readIORef calls `shouldReturn` 1

    it "does not fall back to the next model once the signal fired" $ do
      sig <- newAbortSignal
      (failing, _) <- countingGateway (Left (HttpError 400 "bad"))
      (ok, okCalls) <- countingGateway (Right (ChatResponse "hi" [textPart "hi"] Nothing Nothing))
      let failing' = failing {gwGenerateText = \h req -> abort sig >> failing.gwGenerateText h req}
          models = ModelWithFallbacks (testModelConfig failing') [testModelConfig ok]
      r <- generateTextWithFallbacks (testGenRequest []) {grAbortSignal = Just sig} models
      fmap (.respText) r `shouldBe` Left GErrAborted
      readIORef okCalls `shouldReturn` 0

    it "cuts off a stream once the signal fired" $ do
      sig <- newAbortSignal
      let resp = ChatResponse "abc" [textPart "abc"] Nothing Nothing
          gw =
            (streamGateway [] resp)
              { gwStreamText = \_ _ onEvent -> do
                  onEvent (StreamDelta "a")
                  abort sig
                  onEvent (StreamDelta "b")
                  onEvent (StreamDelta "c")
                  pure (Right resp)
              }
      ref <- newIORef []
      r <- streamTextWithFallbacks (\c -> modifyIORef' ref (c :)) (testGenRequest []) {grAbortSignal = Just sig} (ModelWithFallbacks (testModelConfig gw) [])
      fmap (.respText) r `shouldBe` Left GErrAborted
      reverse <$> readIORef ref `shouldReturn` [TextDelta "a"]

    it "interrupts a stream that stalls after its first event" $ do
      sig <- newAbortSignal
      started <- newEmptyMVar
      closed <- newIORef False
      let resp = ChatResponse "a" [textPart "a"] Nothing Nothing
          gw =
            (streamGateway [] resp)
              { gwStreamText = \_ _ onEvent ->
                  (onEvent (StreamDelta "a") >> putMVar started () >> hang >> pure (Right resp))
                    `finally` writeIORef closed True
              }
          -- The first event arrives in time, so the first-event timeout is off.
          mc = (testModelConfig gw) {mcRequestTimeout = Just 1000}
      abortAfter started sig
      ref <- newIORef []
      r <- timeout 5_000_000 $ streamTextWithFallbacks (\c -> modifyIORef' ref (c :)) (testGenRequest []) {grAbortSignal = Just sig} (ModelWithFallbacks mc [])
      fmap (fmap (.respText)) r `shouldBe` Just (Left GErrAborted)
      readIORef closed `shouldReturn` True
      reverse <$> readIORef ref `shouldReturn` [TextDelta "a"]

    it "interrupts a non-streaming call in flight" $ do
      sig <- newAbortSignal
      started <- newEmptyMVar
      closed <- newIORef False
      let gw =
            (streamGateway [] (ChatResponse "" [] Nothing Nothing))
              { gwGenerateText = \_ _ ->
                  (putMVar started () >> hang >> pure (Left TimeoutError))
                    `finally` writeIORef closed True
              }
      abortAfter started sig
      r <- timeout 5_000_000 $ generateTextWithFallbacks (testGenRequest []) {grAbortSignal = Just sig} (ModelWithFallbacks (testModelConfig gw) [])
      fmap (fmap (.respText)) r `shouldBe` Just (Left GErrAborted)
      readIORef closed `shouldReturn` True

    it "interrupts the throttle sleep before a call" $ do
      sig <- newAbortSignal
      (gw, calls) <- countingGateway (Right (ChatResponse "hi" [textPart "hi"] Nothing Nothing))
      let mc = (testModelConfig gw) {mcThrottleDelay = Just 60_000}
      _ <- forkIO (threadDelay 50_000 >> abort sig)
      r <- timeout 5_000_000 $ generateTextWithFallbacks (testGenRequest []) {grAbortSignal = Just sig} (ModelWithFallbacks mc [])
      fmap (fmap (.respText)) r `shouldBe` Just (Left GErrAborted)
      readIORef calls `shouldReturn` 0

  describe "throttle" $ do
    it "sleeps the throttle delay before every attempt, retries included" $ do
      stamps <- newIORef []
      let gw =
            (streamGateway [] (ChatResponse "" [] Nothing Nothing))
              { gwGenerateText = \_ _ -> do
                  t <- getMonotonicTime
                  modifyIORef' stamps (t :)
                  pure (Left (HttpError 503 "overloaded"))
              }
          mc = (testModelConfig gw) {mcThrottleDelay = Just 100, mcRetryCount = 2}
      start <- getMonotonicTime
      _ <- generateTextLLM (testGenRequest []) mc
      ts <- reverse <$> readIORef stamps
      length ts `shouldBe` 3
      -- Backoff is 0 here, so each gap is the throttle alone.
      let gaps = zipWith (-) ts (start : ts)
      all (>= 0.1) gaps `shouldBe` True

  describe "modelRetryPolicy" $ do
    it "caps the backoff delay at 30 s" $ do
      let mc =
            (testModelConfig (streamGateway [] (ChatResponse "" [] Nothing Nothing)))
              { mcRetryCount = 20,
                mcJitterBackoff = 1000
              }
      delays <-
        mapM
          (\n -> getRetryPolicyM (modelRetryPolicy mc) defaultRetryStatus {rsIterNumber = n})
          [10 .. 19]
      all (maybe False (<= 30_000_000)) delays `shouldBe` True

  describe "stream reset between attempts" $ do
    it "emits StreamAttemptReset before a retry when the failed attempt streamed" $ do
      gw <- flakyGateway "abc" "xyz"
      let mc = (testModelConfig gw) {mcRetryCount = 1}
      ref <- newIORef []
      r <- streamTextLLM (\c -> modifyIORef' ref (c :)) (testGenRequest []) mc
      fmap (.respText) r `shouldBe` Right "xyz"
      chunks <- reverse <$> readIORef ref
      chunks
        `shouldBe` [ TextDelta "abc",
                     StreamAttemptReset,
                     TextDelta "xyz",
                     RoundTextRoleCommitted AnswerRole
                   ]

    it "does not emit StreamAttemptReset when the failed attempt streamed nothing" $ do
      gw <- flakyGateway "" "xyz"
      let mc = (testModelConfig gw) {mcRetryCount = 1}
      ref <- newIORef []
      _ <- streamTextLLM (\c -> modifyIORef' ref (c :)) (testGenRequest []) mc
      chunks <- reverse <$> readIORef ref
      chunks `shouldBe` [TextDelta "xyz", RoundTextRoleCommitted AnswerRole]

    it "emits StreamAttemptReset before falling back to the next model" $ do
      let failing =
            (streamGateway [] (ChatResponse "" [] Nothing Nothing))
              { gwStreamText = \_ _ onEvent -> do
                  onEvent (StreamDelta "abc")
                  pure (Left (HttpError 400 "bad"))
              }
          ok = streamGateway [StreamDelta "xyz"] (ChatResponse "xyz" [textPart "xyz"] Nothing Nothing)
          models = ModelWithFallbacks (testModelConfig failing) [testModelConfig ok]
      ref <- newIORef []
      r <- streamTextWithFallbacks (\c -> modifyIORef' ref (c :)) (testGenRequest []) models
      fmap (.respText) r `shouldBe` Right "xyz"
      chunks <- reverse <$> readIORef ref
      chunks
        `shouldBe` [ TextDelta "abc",
                     StreamAttemptReset,
                     TextDelta "xyz",
                     RoundTextRoleCommitted AnswerRole
                   ]

-- | Block for a minute: long enough that only an abort ends it in a test.
hang :: IO ()
hang = threadDelay 60_000_000

-- | Fire @sig@ from another thread once @started@ is filled.
abortAfter :: MVar () -> AbortSignal -> IO ()
abortAfter started sig = () <$ forkIO (takeMVar started >> abort sig)

-- | First call streams @first@ (if non-empty) then fails with a retryable
-- 503; later calls stream @second@ and succeed.
flakyGateway :: Text -> Text -> IO LLMGateway
flakyGateway first second = do
  attempts <- newIORef (0 :: Int)
  let resp = ChatResponse second [textPart second] Nothing Nothing
  pure
    (streamGateway [] resp)
      { gwStreamText = \_ _ onEvent -> do
          n <- atomicModifyIORef' attempts (\k -> (k + 1, k))
          if n == 0
            then do
              when (first /= "") $ onEvent (StreamDelta first)
              pure (Left (HttpError 503 "overloaded"))
            else do
              onEvent (StreamDelta second)
              pure (Right resp)
      }

-- | A gateway whose non-streaming call always returns @result@, with a
-- counter of how often it was called.
countingGateway :: LLMTextResult -> IO (LLMGateway, IORef Int)
countingGateway result = do
  calls <- newIORef 0
  let gw =
        (streamGateway [] (ChatResponse "" [] Nothing Nothing))
          { gwGenerateText = \_ _ -> modifyIORef' calls (+ 1) >> pure result
          }
  pure (gw, calls)

streamGateway :: [StreamEvent] -> ChatResponse -> LLMGateway
streamGateway events resp =
  LLMGateway
    { gwName = "stream-mock",
      gwImageMediaTypes = supportedImageMediaTypes,
      gwImageUrlSupport = True,
      gwGenerateText = \_ _ -> pure (Right resp),
      gwStreamText = \_ _ onEvent -> mapM_ onEvent events >> pure (Right resp),
      gwGenerateObject = \_ _ _ -> pure (Right (object [], Nothing))
    }

runStream :: LLMGateway -> [Turn] -> IO [StreamChunk]
runStream gw turns = do
  ref <- newIORef []
  let onChunk c = modifyIORef' ref (c :)
  _ <- streamTextLLM onChunk (testGenRequest turns) (testModelConfig gw)
  readIORef ref

testGenRequest :: [Turn] -> GenRequest
testGenRequest turns =
  GenRequest
    { grSystemPrompt = Nothing,
      grMessages = turns,
      grTools = [],
      grAbortSignal = Nothing,
      grLLMHooks = llmHooks noHooks,
      grHooks = noHooks
    }

testModelConfig :: LLMGateway -> ModelConfig
testModelConfig gw =
  ModelConfig
    { mcGateway = gw,
      mcModel = "test",
      mcPricing = zeroPricing,
      mcMaxTokens = 256,
      mcTemperature = Nothing,
      mcThinking = Nothing,
      mcCapabilities = defaultModelCapabilities,
      mcRequestTimeout = Nothing,
      mcThrottleDelay = Nothing,
      mcRetryCount = 0,
      mcJitterBackoff = 0
    }

zeroPricing :: PricingInfo
zeroPricing = defaultPricingInfo 0 0