diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,10 @@ +# Changelog + +## Unreleased + +## 0.1.0.0 - 2026-06-13 + +### Added + +- Initial Hackage release of provider-neutral response caching for shikumi. +- Content-addressed cache keys, in-memory and SQLite cache backends, cached LLM interpretation, and response JSON support shared by backend packages. diff --git a/shikumi-cache.cabal b/shikumi-cache.cabal new file mode 100644 --- /dev/null +++ b/shikumi-cache.cabal @@ -0,0 +1,89 @@ +cabal-version: 3.4 +name: shikumi-cache +version: 0.1.0.0 +synopsis: Content-addressed response caching for shikumi (EP-6) +category: AI +description: + The caching subsystem for shikumi: a content-addressed cache key (a BLAKE3 + digest over a canonical request serialization — the contract EP-7's replay + reproduces), a provider-neutral @Cache@ effect, an in-memory backend, and the + @cachedLLM@ memoizing interpreter that wraps EP-1's @LLM@ effect so identical + requests contact the provider once. + +license: BSD-3-Clause +author: Nadeem Bitar +maintainer: nadeem@gmail.com +build-type: Simple +extra-doc-files: CHANGELOG.md + +common common-options + ghc-options: + -Wall -Wcompat -Widentities -Wincomplete-uni-patterns + -Wincomplete-record-updates -Wredundant-constraints + -fhide-source-paths -Wmissing-export-lists -Wpartial-fields + -Wmissing-deriving-strategies + + default-language: GHC2024 + default-extensions: + DeriveAnyClass + DuplicateRecordFields + MultilineStrings + OverloadedLabels + OverloadedStrings + +library + import: common-options + hs-source-dirs: src + exposed-modules: + Shikumi.Cache + Shikumi.Cache.Backend.Memory + Shikumi.Cache.Backend.SQLite + Shikumi.Cache.Key + Shikumi.Cache.ResponseJSON + Shikumi.Cache.Types + + build-depends: + , aeson + , baikai + , base >=4.20 && <5 + , blake3 + , bytestring + , containers + , direct-sqlite + , effectful + , generic-lens + , lens ^>=5.3 + , scientific + , shikumi ^>=0.1.0.0 + , stm + , text ^>=2.1 + , time + , vector + +test-suite shikumi-cache-test + import: common-options + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + ghc-options: -threaded -with-rtsopts=-N + build-depends: + , aeson + , baikai + , base + , bytestring + , containers + , directory + , effectful + , filepath + , generic-lens + , lens + , process + , shikumi ^>=0.1.0.0 + , shikumi-cache ^>=0.1.0.0 + , stm + , tasty + , tasty-hunit + , temporary + , text + , time + , vector diff --git a/src/Shikumi/Cache.hs b/src/Shikumi/Cache.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Cache.hs @@ -0,0 +1,72 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE TypeFamilies #-} + +-- | The cache storage effect and the memoizing interpreter (EP-6). +-- +-- 'Cache' is a small @effectful@ dynamic effect with two operations — +-- 'lookupCache' and 'storeCache' — interpreted by a backend (memory, SQLite, …). +-- The /policy/ is separate: 'cachedLLM' re-interprets EP-1's @LLM@ effect so that +-- 'Shikumi.LLM.complete' is memoized — identical requests contact the provider +-- once. Keeping mechanism (storage) and policy (when to read/write) apart lets a +-- caller use the store directly, or compose the memoizer, independently. +module Shikumi.Cache + ( -- * The storage effect + Cache (..), + lookupCache, + storeCache, + + -- * The memoizing policy + cachedLLM, + + -- * Re-exports + CacheKey (..), + cacheKey, + currentKeyVersion, + CachedResponse (..), + ) +where + +import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>)) +import Effectful.Dispatch.Dynamic (interpose, passthrough, send) +import Shikumi.Cache.Key (CacheKey (..), cacheKey, currentKeyVersion) +import Shikumi.Cache.Types (CachedResponse (..)) +import Shikumi.Effect.Time (Time, getCurrentTime) +import Shikumi.LLM (LLM (..), complete) + +-- | The cache storage effect: look an entry up by key, or store one. +data Cache :: Effect where + LookupCache :: CacheKey -> Cache m (Maybe CachedResponse) + StoreCache :: CacheKey -> CachedResponse -> Cache m () + +type instance DispatchOf Cache = 'Dynamic + +-- | Look up a cached response by key (a backend MISS is 'Nothing'). +lookupCache :: (Cache :> es) => CacheKey -> Eff es (Maybe CachedResponse) +lookupCache = send . LookupCache + +-- | Store a cached response under a key (idempotent upsert at every backend). +storeCache :: (Cache :> es) => CacheKey -> CachedResponse -> Eff es () +storeCache k v = send (StoreCache k v) + +-- | Memoize EP-1's @LLM@ @complete@ through the 'Cache' effect: on a HIT (a key +-- present whose 'keyVersion' matches the live namespace) return the stored +-- response without contacting the provider; on a MISS delegate to the underlying +-- @LLM@ handler, store the result, and return it. A stale-version entry is +-- treated as a MISS (it is overwritten on re-fetch). The streaming op is passed +-- through unchanged — streams are not cached. +cachedLLM :: + (Cache :> es, LLM :> es, Time :> es) => + Eff es a -> + Eff es a +cachedLLM = interpose $ \env -> \case + Complete model ctx opts -> do + let key = cacheKey model ctx opts + hit <- lookupCache key + case hit of + Just cr | keyVersion cr == currentKeyVersion -> pure (response cr) + _ -> do + resp <- complete model ctx opts + now <- getCurrentTime + storeCache key (CachedResponse resp now currentKeyVersion) + pure resp + other -> passthrough env other diff --git a/src/Shikumi/Cache/Backend/Memory.hs b/src/Shikumi/Cache/Backend/Memory.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Cache/Backend/Memory.hs @@ -0,0 +1,34 @@ +-- | The in-memory cache backend (EP-6): a process-local @STM@ map from the hex +-- cache key to a 'CachedResponse'. It stores the Haskell value directly (no +-- serialization), so it needs no JSON instances for the baikai 'Response' graph. +-- Build the store with 'newMemoryCache' and discharge the 'Cache' effect with +-- 'runCacheMemory'. +module Shikumi.Cache.Backend.Memory + ( MemoryCache, + newMemoryCache, + runCacheMemory, + ) +where + +import Control.Concurrent.STM (TVar, modifyTVar', newTVarIO) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import Effectful (Eff, (:>)) +import Effectful.Concurrent (Concurrent) +import Effectful.Concurrent.STM (atomically, readTVarIO) +import Effectful.Dispatch.Dynamic (interpret) +import Shikumi.Cache (Cache (..), CacheKey (unCacheKey), CachedResponse) + +-- | The in-memory store: a 'TVar' over a map keyed by the hex cache key. +type MemoryCache = TVar (Map Text CachedResponse) + +-- | Create a fresh, empty in-memory cache. +newMemoryCache :: IO MemoryCache +newMemoryCache = newTVarIO Map.empty + +-- | Discharge the 'Cache' effect against an in-memory store. +runCacheMemory :: (Concurrent :> es) => MemoryCache -> Eff (Cache : es) a -> Eff es a +runCacheMemory tv = interpret $ \_ -> \case + LookupCache k -> Map.lookup (unCacheKey k) <$> readTVarIO tv + StoreCache k v -> atomically (modifyTVar' tv (Map.insert (unCacheKey k) v)) diff --git a/src/Shikumi/Cache/Backend/SQLite.hs b/src/Shikumi/Cache/Backend/SQLite.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Cache/Backend/SQLite.hs @@ -0,0 +1,122 @@ +-- | The SQLite cache backend (EP-6) — a durable, embedded, single-file store. +-- +-- Unlike the in-memory backend (which holds the Haskell 'CachedResponse' +-- directly), SQLite persists the entry as JSON in a file, so a cache written by +-- one process is read back by a later process: run a program, kill it, restart, +-- and an identical request is served from disk with no provider call. The JSON +-- round-trip for baikai's 'Baikai.Response.Response' graph comes from +-- "Shikumi.Cache.ResponseJSON" (re-exported through 'CachedResponse''s +-- 'Data.Aeson.ToJSON'/'Data.Aeson.FromJSON'). +-- +-- The store is a thin embedded SQLite database (no server) accessed via +-- @direct-sqlite@. A single 'Database.SQLite3.Database' handle is guarded by an +-- 'MVar' so the 'Cache' effect's lookups and stores are serialized (SQLite's +-- default threading mode does not allow concurrent use of one connection). +module Shikumi.Cache.Backend.SQLite + ( SQLiteCache, + openSQLiteCache, + closeSQLiteCache, + withSQLiteCache, + runCacheSQLite, + ) +where + +import Control.Concurrent.MVar (MVar, newMVar, withMVar) +import Control.Exception (bracket) +import Data.Aeson (decodeStrict', encode) +import Data.ByteString.Lazy qualified as BL +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding qualified as TE +import Data.Time.Format.ISO8601 (iso8601Show) +import Database.SQLite3 + ( ColumnIndex (ColumnIndex), + Database, + SQLData (SQLText), + StepResult (Done, Row), + ) +import Database.SQLite3 qualified as SQL +import Effectful (Eff, IOE, liftIO, (:>)) +import Effectful.Dispatch.Dynamic (interpret) +import Shikumi.Cache (Cache (..), CacheKey (unCacheKey), CachedResponse (..)) + +-- | A handle to an open SQLite-backed cache. The 'Database' is behind an 'MVar' +-- so all access through the 'Cache' effect is serialized. +newtype SQLiteCache = SQLiteCache {sqliteDb :: MVar Database} + +-- | The schema. The @key@ is the 64-hex 'CacheKey' (already version-namespaced +-- via the @version@ field baked into the hash). @value@ is the UTF-8 JSON of +-- 'CachedResponse'. @stored_at@ is ISO-8601 for inspection/eviction. +createTableSQL :: Text +createTableSQL = + """ + CREATE TABLE IF NOT EXISTS shikumi_cache + ( key TEXT PRIMARY KEY + , value TEXT NOT NULL + , stored_at TEXT NOT NULL + ); + """ + +-- | SELECT the JSON value for a key. +selectSQL :: Text +selectSQL = + """ + SELECT value FROM shikumi_cache WHERE key = ?; + """ + +-- | Upsert the JSON value for a key. +insertSQL :: Text +insertSQL = + """ + INSERT OR REPLACE INTO shikumi_cache (key, value, stored_at) + VALUES (?, ?, ?); + """ + +-- | Open (creating if needed) a SQLite cache at the given file path, ensuring +-- the schema exists. Pass @\":memory:\"@ for a transient in-process database. +openSQLiteCache :: FilePath -> IO SQLiteCache +openSQLiteCache path = do + db <- SQL.open (T.pack path) + SQL.exec db createTableSQL + SQLiteCache <$> newMVar db + +-- | Close the underlying database handle. +closeSQLiteCache :: SQLiteCache -> IO () +closeSQLiteCache (SQLiteCache mv) = withMVar mv SQL.close + +-- | Open a SQLite cache, run an action with it, and close it (exception-safe). +withSQLiteCache :: FilePath -> (SQLiteCache -> IO a) -> IO a +withSQLiteCache path = bracket (openSQLiteCache path) closeSQLiteCache + +-- | Discharge the 'Cache' effect against the SQLite store. A row that fails to +-- decode is treated as absent (a MISS), so a corrupt or stale-schema entry is +-- simply re-fetched and overwritten — never returned. +runCacheSQLite :: (IOE :> es) => SQLiteCache -> Eff (Cache : es) a -> Eff es a +runCacheSQLite (SQLiteCache mv) = interpret $ \_ -> \case + LookupCache k -> liftIO (withMVar mv (sqliteLookup k)) + StoreCache k v -> liftIO (withMVar mv (\db -> sqliteStore db k v)) + +-- | SELECT the JSON for a key and decode it; 'Nothing' on absent key or decode +-- failure. +sqliteLookup :: CacheKey -> Database -> IO (Maybe CachedResponse) +sqliteLookup k db = + bracket (SQL.prepare db selectSQL) SQL.finalize $ \stmt -> do + SQL.bindText stmt 1 (unCacheKey k) + res <- SQL.step stmt + case res of + Done -> pure Nothing + Row -> do + sqlVal <- SQL.column stmt (ColumnIndex 0) + pure $ case sqlVal of + SQLText t -> decodeStrict' (TE.encodeUtf8 t) + _ -> Nothing + +-- | Upsert (@INSERT OR REPLACE@) the JSON for a key; idempotent. +sqliteStore :: Database -> CacheKey -> CachedResponse -> IO () +sqliteStore db k v = + bracket (SQL.prepare db insertSQL) SQL.finalize $ \stmt -> do + let json = TE.decodeUtf8 (BL.toStrict (encode v)) + stamp = T.pack (iso8601Show (storedAt v)) + SQL.bind stmt [SQLText (unCacheKey k), SQLText json, SQLText stamp] + _ <- SQL.step stmt + pure () diff --git a/src/Shikumi/Cache/Key.hs b/src/Shikumi/Cache/Key.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Cache/Key.hs @@ -0,0 +1,117 @@ +{-# 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)) diff --git a/src/Shikumi/Cache/ResponseJSON.hs b/src/Shikumi/Cache/ResponseJSON.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Cache/ResponseJSON.hs @@ -0,0 +1,86 @@ +{-# OPTIONS_GHC -Wno-orphans #-} + +-- | A faithful JSON round-trip for baikai's 'Response' graph (EP-6). +-- +-- The persistent cache backends (SQLite/Redis/Postgres) store each cached +-- 'Baikai.Response.Response' as JSON and read it back. baikai's 'Response' graph +-- round-trips only /partially/ out of the box: 'Baikai.Model.Model', +-- 'Baikai.Api.Api', 'Baikai.Content.AssistantContent', and +-- 'Baikai.StopReason.StopReason' have both 'ToJSON' and 'FromJSON'; but +-- 'Baikai.Usage.Usage', 'Baikai.Cost.Cost', and 'Baikai.Cost.CostBreakdown' are +-- 'ToJSON'-only, and 'Baikai.Message.AssistantPayload' and 'Response' have no +-- aeson instances at all. +-- +-- This module supplies the missing pieces as __orphan instances__ so a 'Response' +-- can be encoded and decoded losslessly enough for caching. The instances mirror +-- baikai's own encoders exactly: +-- +-- * 'Usage' uses @camelTo2 '_'@ field labels (baikai's @usageOptions@), so the +-- 'FromJSON' reads the same snake_case keys baikai's 'ToJSON' writes. +-- * 'Cost' / 'CostBreakdown' read the @usd@ / @*_usd@ 'Scientific' fields baikai +-- writes (via @fromRationalRepetendUnlimited@) and lift them back to +-- 'Rational' with 'toRational'. This is lossy only for non-terminating +-- repetends; it never affects the typed-output guarantee, which is decoded +-- from the assistant __text__ ('AssistantContent', which round-trips +-- exactly) — cost is metadata. +-- * 'AssistantPayload' / 'Response' use @defaultOptions@ (matching baikai's +-- @deriving anyclass ToJSON@). +-- +-- This module is the __single home__ for these orphans across the whole +-- framework. EP-7's @Shikumi.Trace.ResponseJSON@ re-exports them from here rather +-- than defining its own copy, so no module that imports both the cache and the +-- trace package ever sees duplicate instances. +module Shikumi.Cache.ResponseJSON () where + +import Baikai.Cost (Cost (..), CostBreakdown (..)) +import Baikai.Message (AssistantPayload (..)) +import Baikai.Response (Response (..)) +import Baikai.Usage (Usage (..)) +import Data.Aeson + ( FromJSON (parseJSON), + Key, + Object, + Options (fieldLabelModifier), + ToJSON (toJSON), + camelTo2, + defaultOptions, + genericParseJSON, + genericToJSON, + withObject, + (.:), + ) +import Data.Aeson.Types (Parser) +import Data.Scientific (Scientific) + +-- | baikai's 'Usage' field-label scheme: @camelTo2 '_'@ (snake_case). +usageOptions :: Options +usageOptions = defaultOptions {fieldLabelModifier = camelTo2 '_'} + +-- | Parse a required object field as 'Scientific' and lift it to 'Rational', +-- inverting baikai's @Rational -> Scientific@ cost encoding. +ratField :: Object -> Key -> Parser Rational +ratField o k = toRational <$> (o .: k :: Parser Scientific) + +instance FromJSON Usage where + parseJSON = genericParseJSON usageOptions + +instance FromJSON CostBreakdown where + parseJSON = withObject "CostBreakdown" $ \o -> + CostBreakdown + <$> ratField o "input_usd" + <*> ratField o "output_usd" + <*> ratField o "cached_input_usd" + <*> ratField o "cached_write_usd" + +instance FromJSON Cost where + parseJSON = withObject "Cost" $ \o -> + Cost <$> ratField o "usd" <*> o .: "breakdown" + +instance FromJSON AssistantPayload where + parseJSON = genericParseJSON defaultOptions + +instance ToJSON Response where + toJSON = genericToJSON defaultOptions + +instance FromJSON Response where + parseJSON = genericParseJSON defaultOptions diff --git a/src/Shikumi/Cache/Types.hs b/src/Shikumi/Cache/Types.hs new file mode 100644 --- /dev/null +++ b/src/Shikumi/Cache/Types.hs @@ -0,0 +1,31 @@ +-- | The value stored in the cache (EP-6): a full provider 'Response' plus the +-- operational metadata needed to age and version-guard an entry. +-- +-- The in-memory backend stores this Haskell value directly (no serialization). +-- The persistent backends (SQLite/Redis/Postgres) store its JSON encoding; the +-- @ToJSON@/@FromJSON@ for the baikai 'Response' graph that baikai itself does not +-- ship in full come from "Shikumi.Cache.ResponseJSON" (imported here for their +-- instances). Decoding an entry whose 'keyVersion' does not match the live +-- namespace is treated as a cache MISS, so a version bump never serves a +-- stale-schema row. +module Shikumi.Cache.Types + ( CachedResponse (..), + ) +where + +import Baikai (Response) +import Data.Aeson (FromJSON, ToJSON) +import Data.Text (Text) +import Data.Time.Clock (UTCTime) +import GHC.Generics (Generic) +import Shikumi.Cache.ResponseJSON () + +-- | A cached provider response: the full 'Response' (message, usage, cost), when +-- it was written, and the key-namespace version it was written under. +data CachedResponse = CachedResponse + { response :: !Response, + storedAt :: !UTCTime, + keyVersion :: !Text + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,253 @@ +-- | EP-6 acceptance: the content-addressed cache key (integration point #7, +-- golden-pinned), the in-memory backend round-trip, the SQLite backend +-- round-trip and cross-process restart durability, the @cachedLLM@ memoizer (the +-- headline one-provider-call behaviour) via a counting stub, and +-- versioning/invalidation. The server-backed Redis/Postgres backends are tested +-- in their own packages (@shikumi-cache-redis@/@shikumi-cache-postgres@). +-- +-- The SQLite restart test re-executes this binary as a subprocess in a "write" +-- and a "read" phase against the same temp file (selected by the +-- @SHIKUMI_SQLITE_PHASE@ / @SHIKUMI_SQLITE_FILE@ env vars), proving the entry +-- survives on disk across OS processes — not merely across handles in one +-- process. +module Main (main) where + +import Baikai (Context, Model, Options, Response, user, _Context, _Model, _Options, _Response) +import Control.Lens ((&), (.~)) +import Data.Generics.Labels () +import Data.IORef (IORef, modifyIORef', newIORef, readIORef) +import Data.Text qualified as T +import Data.Time.Clock (UTCTime) +import Data.Vector qualified as V +import Effectful (Eff, IOE, liftIO, runEff, type (:>)) +import Effectful.Concurrent (runConcurrent) +import Effectful.Dispatch.Dynamic (interpret) +import Shikumi.Cache + ( CacheKey (..), + CachedResponse (..), + cacheKey, + cachedLLM, + currentKeyVersion, + lookupCache, + storeCache, + ) +import Shikumi.Cache.Backend.Memory (newMemoryCache, runCacheMemory) +import Shikumi.Cache.Backend.SQLite (runCacheSQLite, withSQLiteCache) +import Shikumi.Cache.Key (canonicalJSON, requestToCanonicalValueVersioned) +import Shikumi.Effect.Time (runTime) +import Shikumi.LLM (LLM (..), complete) +import System.Environment (getEnvironment, getExecutablePath, lookupEnv) +import System.Exit (ExitCode (ExitSuccess), exitFailure, exitSuccess) +import System.FilePath (()) +import System.IO.Temp (withSystemTempDirectory) +import System.Process (CreateProcess (env), proc, readCreateProcessWithExitCode) +import Test.Tasty (TestTree, defaultMain, testGroup) +import Test.Tasty.HUnit (assertBool, testCase, (@?=)) + +-- --------------------------------------------------------------------------- +-- Fixtures +-- --------------------------------------------------------------------------- + +fixModel :: Model +fixModel = _Model & #modelId .~ "claude-sonnet-4-6" & #provider .~ "anthropic" + +fixCtx :: Context +fixCtx = + _Context + & #systemPrompt .~ Just "You are helpful." + & #messages .~ V.singleton (user "ping") + +fixOpts :: Options +fixOpts = _Options & #temperature .~ Just 0.0 & #maxTokens .~ Just 1024 + +-- | The pinned cache key for @(fixModel, fixCtx, fixOpts)@ — the value EP-7 must +-- reproduce. Captured from a first run; any drift in field set, canonical JSON, +-- or hash breaks this test. +pinnedKey :: T.Text +pinnedKey = "30b2015562ec8b5cd4fdb64c7cc671c84f56f80d24891deec6676c521f008113" + +stubResponse :: Response +stubResponse = _Response + +someTime :: UTCTime +someTime = read "2026-06-08 00:00:00 UTC" + +-- | A counting stub interpreter of EP-1's @LLM@: every completion bumps the +-- counter and returns a fixed response. (The streaming op is unused.) +runCountingLLM :: (IOE :> es) => IORef Int -> Response -> Eff (LLM : es) a -> Eff es a +runCountingLLM ref resp = interpret $ \_ -> \case + Complete {} -> liftIO (modifyIORef' ref (+ 1)) >> pure resp + Stream {} -> pure [] + +-- --------------------------------------------------------------------------- +-- Tests +-- --------------------------------------------------------------------------- + +-- | The fixed key/entry the SQLite restart test writes and reads back. Both the +-- write and read subprocess phases derive them identically from the fixtures. +restartKey :: CacheKey +restartKey = cacheKey fixModel fixCtx fixOpts + +restartEntry :: CachedResponse +restartEntry = CachedResponse stubResponse someTime currentKeyVersion + +main :: IO () +main = do + -- The SQLite restart test re-executes this binary with SHIKUMI_SQLITE_PHASE + -- set, so a "write" then a "read" run hit the same on-disk file as distinct + -- OS processes. These short-circuit before tasty runs. + phase <- lookupEnv "SHIKUMI_SQLITE_PHASE" + case phase of + Just "write" -> sqliteWritePhase + Just "read" -> sqliteReadPhase + _ -> + defaultMain $ + testGroup + "shikumi-cache" + [ keyTests, + memoryTests, + sqliteTests, + memoizeTests, + versioningTests + ] + +-- | Subprocess phase: store the fixed entry into the SQLite file named by +-- @SHIKUMI_SQLITE_FILE@, then exit. Nothing lingers in this process's memory. +sqliteWritePhase :: IO () +sqliteWritePhase = do + file <- requireEnv "SHIKUMI_SQLITE_FILE" + withSQLiteCache file $ \c -> + runEff . runCacheSQLite c $ storeCache restartKey restartEntry + exitSuccess + +-- | Subprocess phase: open the same file fresh and read the entry back; exit 0 +-- iff it decodes to exactly the written entry. +sqliteReadPhase :: IO () +sqliteReadPhase = do + file <- requireEnv "SHIKUMI_SQLITE_FILE" + got <- withSQLiteCache file $ \c -> runEff . runCacheSQLite c $ lookupCache restartKey + if got == Just restartEntry then exitSuccess else exitFailure + +requireEnv :: String -> IO String +requireEnv name = + lookupEnv name >>= maybe (error ("missing env var: " <> name)) pure + +keyTests :: TestTree +keyTests = + testGroup + "cache key" + [ testCase "is a 64-char lowercase hex digest" $ do + let CacheKey hex = cacheKey fixModel fixCtx fixOpts + T.length hex @?= 64 + assertBool "all lowercase hex" (T.all (`elem` ("0123456789abcdef" :: String)) hex), + testCase "is deterministic" $ + cacheKey fixModel fixCtx fixOpts @?= cacheKey fixModel fixCtx fixOpts, + testCase "matches the pinned digest (the contract EP-7 reproduces)" $ do + let CacheKey hex = cacheKey fixModel fixCtx fixOpts + hex @?= pinnedKey, + testCase "a different request yields a different key" $ + assertBool + "temperature change must change the key" + (cacheKey fixModel fixCtx fixOpts /= cacheKey fixModel fixCtx (fixOpts & #temperature .~ Just 0.7)) + ] + +memoryTests :: TestTree +memoryTests = + testGroup + "memory backend" + [ testCase "store-then-lookup returns the entry" $ do + tv <- newMemoryCache + let key = cacheKey fixModel fixCtx fixOpts + entry = CachedResponse stubResponse someTime currentKeyVersion + got <- runEff . runConcurrent . runCacheMemory tv $ (storeCache key entry >> lookupCache key) + got @?= Just entry, + testCase "an absent key returns Nothing" $ do + tv <- newMemoryCache + got <- runEff . runConcurrent . runCacheMemory tv $ lookupCache (CacheKey "deadbeef") + got @?= Nothing + ] + +sqliteTests :: TestTree +sqliteTests = + testGroup + "sqlite" + [ testCase "store-then-lookup round-trips a CachedResponse" $ + withSystemTempDirectory "shikumi-sqlite" $ \dir -> do + let file = dir "cache.db" + got <- + withSQLiteCache file $ \c -> + runEff . runCacheSQLite c $ (storeCache restartKey restartEntry >> lookupCache restartKey) + got @?= Just restartEntry, + testCase "an absent key returns Nothing" $ + withSystemTempDirectory "shikumi-sqlite" $ \dir -> do + let file = dir "cache.db" + got <- withSQLiteCache file $ \c -> runEff . runCacheSQLite c $ lookupCache (CacheKey "deadbeef") + got @?= Nothing, + testCase "restart: an entry written by one OS process is read by another from disk" $ + withSystemTempDirectory "shikumi-sqlite" $ \dir -> do + let file = dir "restart.db" + exe <- getExecutablePath + baseEnv <- getEnvironment + let runPhase ph = + readCreateProcessWithExitCode + (proc exe []) {env = Just (baseEnv ++ [("SHIKUMI_SQLITE_PHASE", ph), ("SHIKUMI_SQLITE_FILE", file)])} + "" + (wc, _, werr) <- runPhase "write" + assertBool ("write phase failed: " <> werr) (wc == ExitSuccess) + (rc, _, rerr) <- runPhase "read" + assertBool ("read phase did not find the on-disk entry: " <> rerr) (rc == ExitSuccess) + ] + +memoizeTests :: TestTree +memoizeTests = + testGroup + "memoize" + [ testCase "same request twice contacts the provider once with equal outputs" $ do + tv <- newMemoryCache + ref <- newIORef 0 + (r1, r2) <- + runEff . runConcurrent . runTime . runCacheMemory tv . runCountingLLM ref stubResponse . cachedLLM $ do + a <- complete fixModel fixCtx fixOpts + b <- complete fixModel fixCtx fixOpts + pure (a, b) + n <- readIORef ref + n @?= 1 + r1 @?= r2, + testCase "two different requests contact the provider twice" $ do + tv <- newMemoryCache + ref <- newIORef 0 + _ <- + runEff . runConcurrent . runTime . runCacheMemory tv . runCountingLLM ref stubResponse . cachedLLM $ do + _ <- complete fixModel fixCtx fixOpts + complete fixModel fixCtx (fixOpts & #temperature .~ Just 0.7) + n <- readIORef ref + n @?= 2 + ] + +versioningTests :: TestTree +versioningTests = + testGroup + "versioning" + [ testCase "a current-version entry is a HIT (no provider call)" $ do + tv <- newMemoryCache + ref <- newIORef 0 + let key = cacheKey fixModel fixCtx fixOpts + runEff . runConcurrent . runCacheMemory tv $ storeCache key (CachedResponse stubResponse someTime currentKeyVersion) + _ <- runEff . runConcurrent . runTime . runCacheMemory tv . runCountingLLM ref stubResponse . cachedLLM $ complete fixModel fixCtx fixOpts + n <- readIORef ref + n @?= 0, + testCase "an entry with a foreign keyVersion is ignored (MISS, provider called)" $ do + tv <- newMemoryCache + ref <- newIORef 0 + let key = cacheKey fixModel fixCtx fixOpts + runEff . runConcurrent . runCacheMemory tv $ storeCache key (CachedResponse stubResponse someTime "shikumi-cache/v0") + _ <- runEff . runConcurrent . runTime . runCacheMemory tv . runCountingLLM ref stubResponse . cachedLLM $ complete fixModel fixCtx fixOpts + n <- readIORef ref + n @?= 1, + testCase "bumping the namespace version changes the hashed bytes" $ + assertBool + "v1 and v2 canonical serializations must differ" + ( canonicalJSON (requestToCanonicalValueVersioned "shikumi-cache/v1" fixModel fixCtx fixOpts) + /= canonicalJSON (requestToCanonicalValueVersioned "shikumi-cache/v2" fixModel fixCtx fixOpts) + ) + ]