baikai-claude-0.7.3.0: test/ShapeSpec.hs

module ShapeSpec (tests) where

import Baikai
import Baikai.Models.Generated qualified as Models
import Baikai.Provider.Claude.Internal.Request (describeThinkingFor, mapRequest)
import Baikai.Provider.Claude.Shape (injectSystemCacheControl, streamRequestBody)
import Baikai.Provider.Claude.Transport qualified as Transport
import Control.Lens ((&), (.~), (^.))
import Control.Monad (forM_)
import Data.Aeson (Value (..), (.=))
import Data.Aeson qualified as Aeson
import Data.Aeson.Key qualified as AesonKey
import Data.Aeson.KeyMap qualified as KeyMap
import Data.Map.Strict qualified as Map
import Data.Text qualified as Text
import Data.Vector qualified as Vector
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=))

tests :: TestTree
tests =
  testGroup
    "ShapeSpec"
    [ testCase "server-side fallbacks stay absent from the request" $ do
        value <- shapedBody Models.anthropic_claude_opus_5 emptyContext emptyOptions
        lookupPath ["fallbacks"] value @?= Nothing,
      summaryDisplayTest,
      fastSpeedTests,
      verbatimToolSchemaTest,
      toolChoiceNoneTest,
      toolCacheControlTest,
      toolCacheControlCompatGateTest,
      systemCacheControlTests
    ]

verbatimToolSchemaTest :: TestTree
verbatimToolSchemaTest =
  testCase "tool input_schema is the caller's verbatim JSON Schema" $ do
    let schema =
          Aeson.object
            [ "type" .= ("object" :: Text.Text),
              "$defs"
                .= Aeson.object
                  [ "unit"
                      .= Aeson.object
                        [ "enum" .= (["c", "f"] :: [Text.Text])
                        ]
                  ],
              "additionalProperties" .= False
            ]
        ctx =
          emptyContext
            & #tools
              .~ Vector.singleton
                (emptyTool & #name .~ "weather" & #description .~ "Weather" & #parameters .~ schema)
    value <- shapedBody Models.anthropic_claude_haiku_4_5 ctx emptyOptions
    lookupPath ["tools", "0", "input_schema"] value @?= Just schema

toolChoiceNoneTest :: TestTree
toolChoiceNoneTest =
  testCase "ToolChoiceNone keeps tools and sends tool_choice none" $ do
    let ctx =
          emptyContext
            & #tools
              .~ Vector.singleton
                (emptyTool & #name .~ "lookup" & #description .~ "Lookup" & #parameters .~ objectSchema)
        opts = emptyOptions & #toolChoice .~ Just ToolChoiceNone
    value <- shapedBody Models.anthropic_claude_haiku_4_5 ctx opts
    assertBool "tools should remain present" (hasNonEmptyTools value)
    lookupPath ["tool_choice", "type"] value @?= Just (String "none")

toolCacheControlTest :: TestTree
toolCacheControlTest =
  testCase "cache marker lands on the last tool definition with ttl" $ do
    let ctx =
          emptyContext
            & #tools
              .~ Vector.fromList
                [ emptyTool & #name .~ "first" & #description .~ "First" & #parameters .~ objectSchema,
                  emptyTool & #name .~ "second" & #description .~ "Second" & #parameters .~ objectSchema
                ]
        opts = emptyOptions & #cacheRetention .~ Just CacheRetentionLong
    value <- shapedBody Models.anthropic_claude_haiku_4_5 ctx opts
    lookupPath ["tools", "0", "cache_control"] value @?= Nothing
    lookupPath ["tools", "1", "cache_control"] value
      @?= Just
        ( Aeson.object
            [ "type" .= ("ephemeral" :: Text.Text),
              "ttl" .= ("1h" :: Text.Text)
            ]
        )

toolCacheControlCompatGateTest :: TestTree
toolCacheControlCompatGateTest =
  testCase "supportsCacheControlOnTools gates tool cache markers" $ do
    let compat =
          defaultAnthropicMessagesCompat
            { supportsCacheControlOnTools = False
            }
        model =
          Models.anthropic_claude_haiku_4_5
            & #compat .~ CompatAnthropicMessages compat
        ctx =
          emptyContext
            & #tools
              .~ Vector.singleton
                (emptyTool & #name .~ "lookup" & #description .~ "Lookup" & #parameters .~ objectSchema)
        opts = emptyOptions & #cacheRetention .~ Just CacheRetentionShort
    value <- shapedBody model ctx opts
    lookupPath ["tools", "0", "cache_control"] value @?= Nothing

systemCacheControlTests :: TestTree
systemCacheControlTests =
  testGroup
    "system prompt cache breakpoint"
    [ testCase "without retention the system prompt stays a plain string" $ do
        value <- shapedBody Models.anthropic_claude_haiku_4_5 systemCtx emptyOptions
        lookupPath ["system"] value @?= Just (String systemText)
        lookupPath ["cache_control"] value @?= Nothing,
      testCase "CacheRetentionNone keeps the system prompt a plain string" $ do
        value <- shapedBody Models.anthropic_claude_haiku_4_5 systemCtx (retention CacheRetentionNone)
        lookupPath ["system"] value @?= Just (String systemText)
        lookupPath ["cache_control"] value @?= Nothing,
      testCase "short retention sends the system prompt as one ephemeral block" $ do
        value <- shapedBody Models.anthropic_claude_haiku_4_5 systemCtx (retention CacheRetentionShort)
        lookupPath ["system"] value @?= Just (systemBlocks ephemeral)
        lookupPath ["cache_control"] value @?= Just ephemeral,
      testCase "long retention gives the system block the 1h ttl" $ do
        value <- shapedBody Models.anthropic_claude_haiku_4_5 systemCtx (retention CacheRetentionLong)
        lookupPath ["system"] value @?= Just (systemBlocks oneHour)
        lookupPath ["cache_control"] value @?= Just oneHour,
      testCase "long retention downgrades with the host" $ do
        let model =
              Models.anthropic_claude_haiku_4_5
                & #compat .~ CompatAnthropicMessages (defaultAnthropicMessagesCompat & #supportsLongCacheRetention .~ False)
        value <- shapedBody model systemCtx (retention CacheRetentionLong)
        lookupPath ["system"] value @?= Just (systemBlocks ephemeral)
        lookupPath ["cache_control"] value @?= Just ephemeral,
      testCase "without a system prompt no system field appears" $ do
        value <- shapedBody Models.anthropic_claude_haiku_4_5 emptyContext (retention CacheRetentionShort)
        lookupPath ["system"] value @?= Nothing
        lookupPath ["cache_control"] value @?= Just ephemeral,
      testCase "tools, system and automatic markers stay within four breakpoints" $ do
        let ctx =
              systemCtx
                & #tools
                  .~ Vector.singleton
                    (emptyTool & #name .~ "lookup" & #description .~ "Lookup" & #parameters .~ objectSchema)
        value <- shapedBody Models.anthropic_claude_haiku_4_5 ctx (retention CacheRetentionShort)
        lookupPath ["tools", "0", "cache_control"] value @?= Just ephemeral
        lookupPath ["system", "0", "cache_control"] value @?= Just ephemeral
        lookupPath ["cache_control"] value @?= Just ephemeral,
      testCase "empty system text never becomes an empty cacheable block" $
        forM_ [CacheRetentionShort, CacheRetentionLong] $ \r -> do
          value <- shapedBody Models.anthropic_claude_haiku_4_5 (emptyContext & #systemPrompt .~ Just "") (retention r)
          lookupPath ["system"] value @?= Just (String ""),
      testCase "system caching is independent of tool-marker support" $ do
        let model =
              Models.anthropic_claude_haiku_4_5
                & #compat .~ CompatAnthropicMessages (defaultAnthropicMessagesCompat & #supportsCacheControlOnTools .~ False)
            ctx = systemCtx & #tools .~ Vector.singleton (emptyTool & #name .~ "lookup" & #parameters .~ objectSchema)
        value <- shapedBody model ctx (retention CacheRetentionLong)
        lookupPath ["tools", "0", "cache_control"] value @?= Nothing
        lookupPath ["system"] value @?= Just (systemBlocks oneHour)
        lookupPath ["cache_control"] value @?= Just oneHour,
      testCase "all markers share the effective TTL with multiple tools" $
        forM_ [True, False] $ \supportsLong -> do
          let model =
                Models.anthropic_claude_haiku_4_5
                  & #compat .~ CompatAnthropicMessages (defaultAnthropicMessagesCompat & #supportsLongCacheRetention .~ supportsLong)
              ctx = systemCtx & #tools .~ Vector.fromList [emptyTool & #name .~ "first", emptyTool & #name .~ "last"]
              marker = if supportsLong then oneHour else ephemeral
          value <- shapedBody model ctx (retention CacheRetentionLong)
          lookupPath ["tools", "0", "cache_control"] value @?= Nothing
          lookupPath ["tools", "1", "cache_control"] value @?= Just marker
          lookupPath ["system"] value @?= Just (systemBlocks marker)
          lookupPath ["cache_control"] value @?= Just marker,
      testCase "different first questions retain the same cacheable system prefix" $ do
        let ctxFor question = systemCtx & #messages .~ Vector.singleton (user question)
        first <- shapedBody Models.anthropic_claude_haiku_4_5 (ctxFor "First question") (retention CacheRetentionShort)
        second <- shapedBody Models.anthropic_claude_haiku_4_5 (ctxFor "Different question") (retention CacheRetentionShort)
        lookupPath ["system"] first @?= lookupPath ["system"] second
        assertBool "user messages remain distinct" (lookupPath ["messages"] first /= lookupPath ["messages"] second),
      testCase "system shaping preserves other fields and existing block arrays" $ do
        let body = Aeson.object ["system" .= systemText, "unrelated" .= objectSchema]
            expected = Aeson.object ["system" .= systemBlocks ephemeral, "unrelated" .= objectSchema]
            inject = injectSystemCacheControl defaultAnthropicMessagesCompat (retention CacheRetentionShort)
        inject body @?= expected
        inject expected @?= expected
    ]
  where
    systemText = "You are a careful assistant." :: Text.Text
    systemCtx = emptyContext & #systemPrompt .~ Just systemText
    retention r = emptyOptions & #cacheRetention .~ Just r
    ephemeral = Aeson.object ["type" .= ("ephemeral" :: Text.Text)]
    oneHour = Aeson.object ["type" .= ("ephemeral" :: Text.Text), "ttl" .= ("1h" :: Text.Text)]
    systemBlocks marker =
      Aeson.toJSON
        [ Aeson.object
            [ "type" .= ("text" :: Text.Text),
              "text" .= systemText,
              "cache_control" .= marker
            ]
        ]

shapedBody :: Model -> Context -> Options -> IO Value
shapedBody model ctx opts = do
  (req, _) <- either (assertFailure . Text.unpack) pure (mapRequest model ctx opts)
  pure (streamRequestBody (anthropicMessagesCompatFor model) ctx opts req)

objectSchema :: Value
objectSchema = Aeson.object ["type" .= ("object" :: Text.Text)]

hasNonEmptyTools :: Value -> Bool
hasNonEmptyTools value =
  case lookupPath ["tools"] value of
    Just (Array xs) -> not (Vector.null xs)
    _ -> False

lookupPath :: [Text.Text] -> Value -> Maybe Value
lookupPath [] value = Just value
lookupPath (field : rest) (Object obj) =
  KeyMap.lookup (AesonKey.fromText field) obj >>= lookupPath rest
lookupPath (field : rest) (Array xs)
  | [(i, "")] <- reads (Text.unpack field),
    i >= 0,
    i < Vector.length xs =
      lookupPath rest (xs Vector.! i)
lookupPath _ _ = Nothing

fastSpeedTests :: TestTree
fastSpeedTests =
  testGroup
    "speed"
    [ testCase "a fast-mode request on claude-opus-5 carries speed and the beta header" $ do
        let m = Models.anthropic_claude_opus_5
            opts = emptyOptions & #speed .~ Just SpeedFast
        value <- shapedBody m emptyContext opts
        lookupPath ["speed"] value @?= Just (String "fast")
        lookup "anthropic-beta" (headers m opts) @?= Just "fast-mode-2026-02-01",
      testCase "a fast-mode request on claude-sonnet-5 omits speed and records the drop" $ do
        let m = Models.anthropic_claude_sonnet_5
            opts = emptyOptions & #speed .~ Just SpeedFast
        value <- shapedBody m emptyContext opts
        lookupPath ["speed"] value @?= Nothing
        lookup "anthropic-beta" (headers m opts) @?= Nothing
        case mapRequest m emptyContext opts of
          Left err -> assertFailure (show err)
          Right (_, translation) -> do
            translation @?= describeThinkingFor m opts
            translation ^. #adjustments @?= [FastModeDroppedUnsupportedModel]
        weakensThinking FastModeDroppedUnsupportedModel @?= False
        Aeson.fromJSON (Aeson.toJSON FastModeDroppedUnsupportedModel) @?= Aeson.Success FastModeDroppedUnsupportedModel,
      testCase "absent and explicit standard speed remain distinct on unsupported models" $ do
        let m = Models.anthropic_claude_sonnet_5
            opts = emptyOptions & #speed .~ Just SpeedStandard
        absent <- shapedBody m emptyContext emptyOptions
        standard <- shapedBody m emptyContext opts
        lookupPath ["speed"] absent @?= Nothing
        lookupPath ["speed"] standard @?= Just (String "standard")
        lookup "anthropic-beta" (headers m opts) @?= Nothing
        describeThinkingFor m opts ^. #adjustments @?= [],
      testCase "caller beta headers override model and automatic fast beta headers" $ do
        let m = Models.anthropic_claude_opus_5 & #headers .~ Map.singleton "anthropic-beta" "model-beta"
            opts = emptyOptions & #speed .~ Just SpeedFast & #headers .~ Map.singleton "Anthropic-Beta" "caller-beta"
        lookup "anthropic-beta" (headers m emptyOptions) @?= Just "model-beta"
        lookup "anthropic-beta" (headers m opts) @?= Just "caller-beta"
    ]
  where
    headers m opts = Transport.requestHeaders "test-key" Nothing (anthropicMessagesCompatFor m) emptyContext m opts

summaryDisplayTest :: TestTree
summaryDisplayTest = testCase "adaptive requests ask for summarized display; budget and absent thinking retain their shapes" $ do
  let opts = emptyOptions & #thinking .~ Just ThinkingLow
  forM_ [Models.anthropic_claude_opus_5, Models.anthropic_claude_opus_4_6, Models.anthropic_claude_fable_5_1, Models.anthropic_claude_opus_5 & #modelId .~ "renamed"] $ \m -> do
    body <- shapedBody m emptyContext opts
    lookupPath ["thinking"] body @?= Just (Aeson.object ["type" .= ("adaptive" :: Text.Text), "display" .= ("summarized" :: Text.Text)])
    describeThinkingFor m opts ^. #displayText @?= Just "summarized"
    absent <- shapedBody m emptyContext emptyOptions
    lookupPath ["thinking"] absent @?= Nothing
    describeThinkingFor m emptyOptions ^. #displayText @?= Nothing
  let budget = Models.anthropic_claude_haiku_4_5
  body <- shapedBody budget emptyContext opts
  lookupPath ["thinking"] body @?= Just (Aeson.object ["type" .= ("enabled" :: Text.Text), "budget_tokens" .= thinkingTokenBudget ThinkingLow])
  describeThinkingFor budget opts ^. #displayText @?= Nothing