Key.hs

{-# LANGUAGE DataKinds #-}

-- | The content-addressed cache key (EP-6) — the MasterPlan's integration point
-- #7. A 'CacheKey' is a BLAKE3 256-bit digest (64 lowercase hex chars) over a
-- /canonical/ JSON serialization of everything about a request that can change
-- the model's answer: the model routing identity, the rendered prompt, the
-- tools, and the per-call options. Latency, response ids, API keys, timeouts,
-- headers, and per-call metadata are excluded — they do not affect the output.
--
-- The serialization is canonical (object keys sorted by Unicode code point, no
-- insignificant whitespace, UTF-8) so the bytes — and therefore the key — are
-- reproducible. The replay engine in
-- @docs/plans/7-hierarchical-tracing-observability-and-replay.md@ (EP-7) reuses
-- 'cacheKey' verbatim, so both plans agree byte-for-byte; the golden test in the
-- test suite pins the exact hex for a fixed request to catch any drift.
module Shikumi.Cache.Key
  ( CacheKey (..),
    cacheKey,
    canonicalJSON,
    requestToCanonicalValue,
    requestToCanonicalValueVersioned,
    currentKeyVersion,
  )
where

import BLAKE3 qualified as B
import Baikai (Context, Model, Options)
import Control.Lens ((^.))
import Data.Aeson (Value (Array, Object, String), object, toJSON, (.=))
import Data.Aeson.Encoding (fromEncoding)
import Data.Aeson.Key qualified as Key
import Data.Aeson.KeyMap qualified as KM
import Data.Aeson.Types (toEncoding)
import Data.ByteString (ByteString)
import Data.ByteString.Builder qualified as BB
import Data.ByteString.Lazy qualified as BL
import Data.Foldable (toList)
import Data.Generics.Labels ()
import Data.List (intersperse, sortOn)
import Data.Scientific (Scientific)
import Data.Text (Text)
import Data.Text qualified as T

-- | The current cache key-namespace version. Baked into every hashed request
-- (the @version@ field) and into 'Shikumi.Cache.Types.CachedResponse.keyVersion'.
-- Bumping it changes every key, making all prior entries unreachable — a clean
-- invalidation with no row deletion.
currentKeyVersion :: Text
currentKeyVersion = "shikumi-cache/v1"

-- | A content-addressed cache key: the 64-hex BLAKE3 digest of a request's
-- canonical serialization. A @newtype@ so it cannot be confused with arbitrary
-- text.
newtype CacheKey = CacheKey {unCacheKey :: Text}
  deriving stock (Eq, Ord, Show)

-- | The cache key for a request triple @(Model, Context, Options)@. The
-- response-format / JSON-schema slot is read from @Options.responseFormat@
-- (provided by EP-2): a request with no schema and one with a schema hash
-- differently, as they must.
cacheKey :: Model -> Context -> Options -> CacheKey
cacheKey m ctx opts =
  let bytes = BL.toStrict (canonicalJSON (requestToCanonicalValue m ctx opts))
      digest = B.hash Nothing [bytes :: ByteString] :: B.Digest 32
   in CacheKey (T.pack (show digest))

-- | Assemble the canonical request 'Value' for the current namespace version.
requestToCanonicalValue :: Model -> Context -> Options -> Value
requestToCanonicalValue = requestToCanonicalValueVersioned currentKeyVersion

-- | Assemble the canonical request 'Value' under an explicit version string.
-- Exposed so a test can prove that bumping the version changes the key. The
-- top-level object carries exactly the integration-point-#7 field set:
-- @api, maxTokens, messages, model, provider, responseFormat, systemPrompt,
-- temperature, thinking, toolChoice, tools, version@. 'canonicalJSON' sorts the
-- keys, so the listing order here is irrelevant.
requestToCanonicalValueVersioned :: Text -> Model -> Context -> Options -> Value
requestToCanonicalValueVersioned version m ctx opts =
  object
    [ "version" .= version,
      "model" .= (m ^. #modelId),
      "provider" .= (m ^. #provider),
      "api" .= toJSON (m ^. #api),
      "systemPrompt" .= (ctx ^. #systemPrompt),
      "messages" .= toJSON (ctx ^. #messages),
      "tools" .= toJSON (ctx ^. #tools),
      "toolChoice" .= toJSON (opts ^. #toolChoice),
      "temperature" .= (toScientific <$> (opts ^. #temperature)),
      "maxTokens" .= (opts ^. #maxTokens),
      "thinking" .= toJSON (opts ^. #thinking),
      "responseFormat" .= toJSON (opts ^. #responseFormat)
    ]
  where
    -- Normalize the Double through Scientific so two requests that are == in
    -- temperature serialize identically regardless of how the Double was built.
    toScientific :: Double -> Scientific
    toScientific = realToFrac

-- | Render a 'Value' to canonical JSON bytes: every object's keys are emitted in
-- ascending Unicode code-point order, with no insignificant whitespace. We do
-- not trust aeson's @encode@ key order (not contractually stable); we sort keys
-- ourselves and reuse aeson's (stable) scalar/string escaping.
canonicalJSON :: Value -> BL.ByteString
canonicalJSON = BB.toLazyByteString . go
  where
    go (Object o) =
      let kvs = sortOn fst [(Key.toText k, v) | (k, v) <- KM.toList o]
       in BB.char7 '{'
            <> mconcat (intersperse (BB.char7 ',') [encStr k <> BB.char7 ':' <> go v | (k, v) <- kvs])
            <> BB.char7 '}'
    go (Array xs) =
      BB.char7 '['
        <> mconcat (intersperse (BB.char7 ',') (map go (toList xs)))
        <> BB.char7 ']'
    go v = fromEncoding (toEncoding v)
    -- Encode an object key as a JSON string (aeson's escaping is stable).
    encStr t = fromEncoding (toEncoding (String t))