diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,34 @@ +# Changelog + +## 0.1.0.0 — 2026-07-14 + +### Added + +- Added event-sourced memory commands, read models, merge and forget behavior, full-detail reads, + and scoped active-memory lookup. +- Added event-sourced sessions with turns, delegation lineage, awaiting/resume state, scoped reads, + and aggregate-enforced identity and correlation invariants. +- Added hybrid full-text and pgvector recall using Reciprocal Rank Fusion, vector capability and + dimension checks, direct memory lookup, and continuous or one-shot embedding workers. +- Added L1 memory extraction and consolidation, L2 scene regeneration, and L3 persona distillation + through Shikumi programs, durable timers, and worker dispatch. +- Added code-driven read-model registration and reconciliation plus compatibility decoding for + legacy Rei events. + +### Fixed + +- Prevented selective scopes from starving the approximate-nearest-neighbor recall channel and + ordered vector candidates by distance before fusion. +- Made atom and scope identities deterministic and collision-safe, surfaced extraction and + consolidation failures, and validated LLM outputs before accepting them. +- Propagated memory forget and confidence changes through scene and persona artifacts. +- Made session lineage validation cycle-safe and returned conflicts for non-idempotent command + replays instead of reporting false success. +- Classified embedding failures into retry, dead-letter, and halt outcomes and wired recall-based + merge candidates into timer processing. + +### Changed + +- Updated the application environment for Keiro 0.3 and Kiroku Store 0.3: it now carries connection + settings, acquires a `KirokuStoreResource`, and explicitly registers read models at startup. +- Removed the unused `embedBatched` API, whose implementation did not perform batched requests. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2026, Nadeem Bitar +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/kioku-core.cabal b/kioku-core.cabal new file mode 100644 --- /dev/null +++ b/kioku-core.cabal @@ -0,0 +1,164 @@ +cabal-version: 3.0 +name: kioku-core +version: 0.1.0.0 +synopsis: Reusable agent memory runtime +description: + Core runtime for kioku. M1 establishes the application effect stack; later + milestones add the memory and session aggregates. + +license: BSD-3-Clause +license-file: LICENSE +author: Nadeem Bitar +maintainer: nadeem@gmail.com +copyright: 2026 Nadeem Bitar +category: Data +build-type: Simple +tested-with: GHC >=9.12 && <9.13 +homepage: https://github.com/shinzui/kioku +bug-reports: https://github.com/shinzui/kioku/issues +extra-doc-files: CHANGELOG.md + +source-repository head + type: git + location: https://github.com/shinzui/kioku.git + subdir: kioku-core + +common warnings + ghc-options: + -Wall -Wcompat -Widentities -Wincomplete-record-updates + -Wincomplete-uni-patterns -Wpartial-fields -Wredundant-constraints + +common shared + default-language: GHC2024 + default-extensions: + BlockArguments + DeriveAnyClass + DuplicateRecordFields + MultilineStrings + OverloadedLabels + OverloadedRecordDot + OverloadedStrings + QualifiedDo + TemplateHaskell + +library + import: warnings, shared + hs-source-dirs: src + exposed-modules: + Kioku.App + Kioku.Distill.Consolidate + Kioku.Distill.Extract + Kioku.Distill.L1 + Kioku.Distill.L2 + Kioku.Distill.L3 + Kioku.Distill.Persona + Kioku.Distill.Runtime + Kioku.Distill.Scene + Kioku.Distill.ScopeIdentity + Kioku.Distill.Timer + Kioku.Distill.Timer.Outcome + Kioku.Distill.Timer.Worker + Kioku.Memory + Kioku.Memory.Domain + Kioku.Memory.Embedding + Kioku.Memory.Embedding.Worker + Kioku.Memory.EventStream + Kioku.Memory.ReadModel + Kioku.ReadModel + Kioku.Recall + Kioku.Recall.Capability + Kioku.Session + Kioku.Session.Domain + Kioku.Session.EventStream + Kioku.Session.ReadModel + Kioku.Worker.Failure + + build-depends: + , aeson >=2.2 && <2.3 + , baikai ^>=0.3.0.0 + , baikai-claude ^>=0.3.0.0 + , baikai-effectful ^>=0.3.0.0 + , base >=4.21 && <5 + , bytestring >=0.11 && <0.13 + , containers >=0.6 && <0.8 + , contravariant >=1.5 && <1.6 + , contravariant-extras >=0.3 && <0.4 + , crypton ^>=1.1.4 + , directory >=1.3 && <1.4 + , effectful >=2.5 && <2.7 + , effectful-core >=2.5 && <2.7 + , filepath >=1.4 && <1.6 + , generic-lens >=2.2 && <2.4 + , hasql >=1.6 && <1.11 + , hasql-pool >=1.2 && <1.5 + , hasql-transaction >=1.0 && <1.3 + , hs-opentelemetry-api >=1.0 && <1.1 + , keiki ^>=0.2.0.0 + , keiro ^>=0.3.0.0 + , keiro-core ^>=0.3.0.0 + , kioku-api ^>=0.1.0.0 + , kiroku-store ^>=0.3.0.1 + , lens >=5.2 && <5.4 + , mmzk-typeid >=0.7 && <0.8 + , shibuya-core >=0.8.0.1 && <0.9 + , shibuya-kiroku-adapter ^>=0.4.0.0 + , shikumi ^>=0.3.0.0 + , shikumi-trace ^>=0.2.0.0 + , text >=2.1 && <2.2 + , time >=1.12 && <1.15 + , uuid >=1.3 && <1.4 + , vector ^>=0.13.2.0 + +test-suite kioku-test + import: warnings, shared + type: exitcode-stdio-1.0 + main-is: Main.hs + other-modules: + Kioku.AwaitingSpec + Kioku.DistillSpec + Kioku.EmbeddingWorkerSpec + Kioku.IdempotencySpec + Kioku.ReadModelReconcileSpec + Kioku.RecallHarness + Kioku.RecallSpec + Kioku.RecallSqlSpec + Kioku.ReiCompatSpec + Kioku.SchemaSpec + Kioku.ScopeIdentitySpec + Kioku.SessionInvariantsSpec + Kioku.SessionLineageSpec + Kioku.TimerWorkerSpec + + hs-source-dirs: test + ghc-options: -threaded -rtsopts -with-rtsopts=-N + build-depends: + , aeson >=2.2 + , baikai ^>=0.3.0.0 + , base >=4.21 && <5 + , bytestring >=0.11 + , containers >=0.6 + , contravariant >=1.5 + , directory >=1.3 + , effectful >=2.5 + , effectful-core >=2.5 + , hasql >=1.6 + , hasql-transaction >=1.0 + , keiro ^>=0.3.0.0 + , keiro-core ^>=0.3.0.0 + , kioku-api ^>=0.1.0.0 + , kioku-core ^>=0.1.0.0 + , kioku-migrations:test-support ^>=0.1.0.0 + , kiroku-store ^>=0.3.0.1 + , lens >=5.2 + , shibuya-core >=0.8.0.1 && <0.9 + , shikumi ^>=0.3.0.0 + , shikumi-trace ^>=0.2.0.0 + , tasty >=1.5 + , tasty-expected-failure >=0.12 + , tasty-hunit >=0.10 + , temporary >=1.3 + , text >=2.1 + , time >=1.12 + , unordered-containers >=0.2 + , uuid >=1.3 + , vector diff --git a/src/Kioku/App.hs b/src/Kioku/App.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/App.hs @@ -0,0 +1,63 @@ +{-# LANGUAGE DataKinds #-} + +module Kioku.App + ( AppEffects, + AppEnv (..), + KeiroMetrics, + runAppIO, + withNoopAppEnv, + noopTracer, + ) +where + +import Effectful (Eff, IOE, runEff) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Keiro.Telemetry (KeiroMetrics) +import Kioku.Prelude +import Kioku.ReadModel (registerKiokuReadModels) +import Kiroku.Store.Connection (ConnectionSettings) +import Kiroku.Store.Effect (Store, runStoreResource) +import Kiroku.Store.Effect.Resource (KirokuStoreResource, withKirokuStore) +import Kiroku.Store.Error (StoreError) +import OpenTelemetry.Attributes qualified as Attr +import OpenTelemetry.Trace.Core qualified as OTel +import Shibuya.Telemetry.Effect (Tracer, Tracing, runTracing) + +type AppEffects = '[Store, KirokuStoreResource, Error StoreError, Tracing, IOE] + +data AppEnv = AppEnv + { connectionSettings :: !ConnectionSettings, + tracer :: !Tracer, + metrics :: !(Maybe KeiroMetrics) + } + deriving stock (Generic) + +runAppIO :: AppEnv -> Eff AppEffects a -> IO (Either StoreError a) +runAppIO env = + runEff + . runTracing (tracer env) + . runErrorNoCallStack + . withKirokuStore (connectionSettings env) + . runStoreResource + +withNoopAppEnv :: ConnectionSettings -> (AppEnv -> IO a) -> IO a +withNoopAppEnv connectionSettings continue = do + tracer <- noopTracer + let env = AppEnv {connectionSettings, tracer, metrics = Nothing} + registration <- runAppIO env registerKiokuReadModels + case registration of + Left err -> fail ("Kioku read-model registration failed: " <> show err) + Right () -> continue env + +noopTracer :: IO Tracer +noopTracer = do + provider <- OTel.createTracerProvider [] OTel.emptyTracerProviderOptions + pure (OTel.makeTracer provider instrumentationLib OTel.tracerOptions) + where + instrumentationLib = + OTel.InstrumentationLibrary + { OTel.libraryName = "kioku-noop", + OTel.libraryVersion = "", + OTel.librarySchemaUrl = "", + OTel.libraryAttributes = Attr.emptyAttributes + } diff --git a/src/Kioku/Distill/Consolidate.hs b/src/Kioku/Distill/Consolidate.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/Consolidate.hs @@ -0,0 +1,108 @@ +{-# LANGUAGE DataKinds #-} + +-- | Pure shikumi program for deciding how an extracted atom affects active memory. +module Kioku.Distill.Consolidate + ( ConsolidationAction (..), + ConsolidationDecision (..), + ConsolidateInput (..), + ExistingMemory (..), + consolidateSignature, + consolidateProgram, + ) +where + +import Kioku.Distill.Extract (ExtractedAtom) +import Kioku.Id (MemoryId, parseIdLenient) +import Kioku.Prelude +import Shikumi.Adapter (ToPrompt) +import Shikumi.Module (predict) +import Shikumi.Program (Program) +import Shikumi.Schema (FromModel, ToSchema, Validatable (..)) +import Shikumi.Schema.Types (Field) +import Shikumi.Signature (Signature, mkSignature) + +data ConsolidationAction + = StoreAtom + | UpdateAtom + | MergeAtom + | SkipAtom + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel) + +data ExistingMemory = ExistingMemory + { memoryId :: Field "existing memory identifier" Text, + memoryType :: Field "existing memory type or category" Text, + content :: Field "existing memory content" Text, + priority :: Field "existing memory priority where lower is more important" Int, + confidence :: Field "existing confidence label" Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +data ConsolidateInput = ConsolidateInput + { scopeLabel :: Field "human-readable memory scope label" Text, + candidate :: ExtractedAtom, + existing :: [ExistingMemory] + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +data ConsolidationDecision = ConsolidationDecision + { action :: ConsolidationAction, + targetMemoryIds :: [Text], + resultContent :: Maybe Text, + rationale :: Field "one concise reason for the decision" Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +-- | An @UpdateAtom@ or @MergeAtom@ that names no parseable target is a decision +-- L1 can only degrade to a store; rejecting it here means the audit trail can +-- never record a merge that did not happen. @StoreAtom@ and @SkipAtom@ have a +-- safe canonical form — no targets — so stray ids are cleared rather than +-- rejected. +instance Validatable ConsolidationDecision where + validate decision = + case decision.action of + StoreAtom -> Right decision {targetMemoryIds = []} + SkipAtom -> Right decision {targetMemoryIds = []} + UpdateAtom -> requireTargets decision + MergeAtom -> requireTargets decision + +requireTargets :: ConsolidationDecision -> Either Text ConsolidationDecision +requireTargets decision + | null decision.targetMemoryIds = + Left (actionLabel decision.action <> " requires at least one targetMemoryId") + | otherwise = do + _ <- traverse parseTarget decision.targetMemoryIds + Right decision + where + parseTarget :: Text -> Either Text MemoryId + parseTarget raw = + case parseIdLenient raw of + Left err -> Left ("targetMemoryIds contains an unparseable id " <> raw <> ": " <> err) + Right mid -> Right mid + +actionLabel :: ConsolidationAction -> Text +actionLabel = \case + StoreAtom -> "StoreAtom" + UpdateAtom -> "UpdateAtom" + MergeAtom -> "MergeAtom" + SkipAtom -> "SkipAtom" + +consolidateSignature :: Signature ConsolidateInput ConsolidationDecision +consolidateSignature = + mkSignature + "Decide how one extracted memory atom should affect the active memories in \ + \the same scope. Choose StoreAtom when the candidate is durable and not \ + \already represented. Choose UpdateAtom when one existing memory should be \ + \rewritten with fresher or clearer content. Choose MergeAtom when multiple \ + \existing memories should collapse into one memory; put all affected ids in \ + \targetMemoryIds and put the winning merged content in resultContent. Choose \ + \SkipAtom when the candidate is transient, unsupported, redundant, or unsafe \ + \to store. For StoreAtom, leave targetMemoryIds empty and put the candidate \ + \content in resultContent. For SkipAtom, leave targetMemoryIds empty and \ + \resultContent absent." + +consolidateProgram :: Program ConsolidateInput ConsolidationDecision +consolidateProgram = predict consolidateSignature diff --git a/src/Kioku/Distill/Extract.hs b/src/Kioku/Distill/Extract.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/Extract.hs @@ -0,0 +1,105 @@ +{-# LANGUAGE DataKinds #-} + +-- | Pure shikumi program for extracting L1 memory atoms from recent session text. +module Kioku.Distill.Extract + ( ExtractInput (..), + ExtractedAtom (..), + ExtractOutput (..), + extractSignature, + extractProgram, + ) +where + +import Data.Text qualified as Text +import Kioku.Prelude +import Shikumi.Adapter (ToPrompt) +import Shikumi.Module (predict) +import Shikumi.Program (Program) +import Shikumi.Schema (FromModel, ToSchema, Validatable (..)) +import Shikumi.Schema.Types (Field, field, unField) +import Shikumi.Signature (Signature, mkSignature) + +data ExtractInput = ExtractInput + { focus :: Field "the session focus, task, or topic" Text, + scopeLabel :: Field "human-readable memory scope label" Text, + conversation :: Field "recent turns, notes, or recorded memory evidence" Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +data ExtractedAtom = ExtractedAtom + { atomType :: Field "one of: fact | pattern | preference | constraint | instruction" Text, + content :: Field "one concise durable memory sentence" Text, + priority :: Field "0=always inject; 100=default; larger=lower priority" Int, + confidence :: Field "one of: high | medium | low" Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +newtype ExtractOutput = ExtractOutput + { atoms :: [ExtractedAtom] + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +atomTypes :: [Text] +atomTypes = ["fact", "pattern", "preference", "constraint", "instruction"] + +confidences :: [Text] +confidences = ["high", "medium", "low"] + +-- | Normalize what has a safe canonical form; reject what does not. +-- +-- An unclamped negative priority is the dangerous case: @priority@ flows +-- straight into @RecordMemoryData@ and from there into every +-- @ORDER BY priority ASC@ read, so a hallucinated @-1000000@ would dominate +-- every candidate and injection query in the scope forever. Unknown enum values +-- were previously coerced to defaults inside L1, silently relabelling a memory; +-- rejecting them surfaces as a shikumi 'Shikumi.Error.ValidationFailure', i.e. a +-- failed extraction, i.e. a retryable timer fire. +instance Validatable ExtractOutput where + validate output = ExtractOutput <$> traverse validateAtom output.atoms + +validateAtom :: ExtractedAtom -> Either Text ExtractedAtom +validateAtom atom = do + content <- nonEmpty "atom content" (Text.strip (unField atom.content)) + atomType <- oneOf "atomType" atomTypes (normalize (unField atom.atomType)) + confidence <- oneOf "confidence" confidences (normalize (unField atom.confidence)) + pure + ExtractedAtom + { atomType = field atomType, + content = field content, + priority = field (clamp 0 100 (unField atom.priority)), + confidence = field confidence + } + where + normalize = Text.toLower . Text.strip + +nonEmpty :: Text -> Text -> Either Text Text +nonEmpty label value + | Text.null value = Left (label <> " must be non-empty") + | otherwise = Right value + +oneOf :: Text -> [Text] -> Text -> Either Text Text +oneOf label allowed value + | value `elem` allowed = Right value + | otherwise = + Left (label <> " must be one of " <> Text.intercalate " | " allowed <> "; got: " <> value) + +clamp :: Int -> Int -> Int -> Int +clamp lo hi = max lo . min hi + +extractSignature :: Signature ExtractInput ExtractOutput +extractSignature = + mkSignature + "Extract durable L1 memory atoms from the provided session text. Include only \ + \information that is useful after this session: stable user preferences, \ + \project facts, recurring patterns, constraints, and explicit instructions. \ + \Ignore transient chatter, one-off status updates, secrets, and unsupported \ + \guesses. Keep each atom as one concise sentence, choose an atomType from the \ + \allowed set, set priority lower for more important memories, and set \ + \confidence to high, medium, or low. Return an empty atoms list when there is \ + \nothing durable to retain." + +extractProgram :: Program ExtractInput ExtractOutput +extractProgram = predict extractSignature diff --git a/src/Kioku/Distill/L1.hs b/src/Kioku/Distill/L1.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/L1.hs @@ -0,0 +1,665 @@ +{-# LANGUAGE DataKinds #-} + +module Kioku.Distill.L1 + ( FindMergeCandidates (..), + L1Error (..), + L1Outcome (..), + L1RunMode (..), + L1Summary (..), + distillSessionL1, + recallCandidates, + scopedScanCandidates, + ) +where + +import Baikai.Embedding (EmbeddingModel) +import Control.Monad (foldM) +import Data.Aeson qualified as Aeson +import Data.ByteString qualified as BS +import Data.ByteString.Lazy qualified as BL +import Data.Either (lefts) +import Data.Functor.Contravariant ((>$<)) +import Data.Int (Int32) +import Data.KindID.V7 qualified as KindID +import Data.List (nub) +import Data.Maybe (catMaybes) +import Data.Set qualified as Set +import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE +import Data.UUID (UUID) +import Data.UUID qualified as UUID +import Data.UUID.V5 qualified as UUIDv5 +import Effectful (Eff, IOE, (:>)) +import Effectful.Error.Static (Error) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Keiro.ReadModel (ReadModelError) +import Kioku.Api.Scope (MemoryScope, scopeFromColumns, scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Api.Types (Confidence (..), MemoryRecord (..), MemoryType (..), confidenceFromText, memoryTypeFromText) +import Kioku.Distill.Consolidate + ( ConsolidateInput (..), + ConsolidationAction (..), + ConsolidationDecision (..), + ExistingMemory (..), + ) +import Kioku.Distill.Extract (ExtractInput (..), ExtractOutput (..), ExtractedAtom (..)) +import Kioku.Distill.Runtime (DistillRuntime, runConsolidation, runExtraction) +import Kioku.Id (MemoryId, SessionId, idText, parseIdLenient) +import Kioku.Memory qualified as Memory +import Kioku.Memory.Domain (RecordMemoryData (..)) +import Kioku.Memory.ReadModel (MemoryRow (..)) +import Kioku.Prelude +import Kioku.Recall qualified as Recall +import Kioku.Recall.Capability (VectorCapability) +import Kioku.Session qualified as Session +import Kioku.Session.ReadModel (SessionRow (..), TurnRow (..)) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Transaction (runTransaction) +import Shikumi.Schema.Types (field, unField) + +newtype FindMergeCandidates es = FindMergeCandidates + { runFindMergeCandidates :: MemoryScope -> Text -> Eff es (Either ReadModelError [MemoryRecord]) + } + +data L1Error + = L1SessionReadFailed !ReadModelError + | L1SessionNotFound !SessionId + | L1TurnReadFailed !ReadModelError + | L1MemoryReadFailed !ReadModelError + | L1ExtractionFailed !Text + | L1ConsolidationFailed !Text + | L1MemoryWriteFailed !Memory.MemoryWriteError + deriving stock (Generic, Show) + +data L1Summary = L1Summary + { extracted :: !Int, + stored :: !Int, + merged :: !Int, + skipped :: !Int + } + deriving stock (Generic, Eq, Show) + +-- | Whether a pass may skip itself when the per-session watermark shows no new +-- turns since the last fully successful pass. Timer fires use +-- 'RespectWatermark' (the debounce); @kioku distill --force@ uses +-- 'IgnoreWatermark'. +data L1RunMode + = RespectWatermark + | IgnoreWatermark + deriving stock (Generic, Eq, Show) + +data L1Outcome + = L1Distilled !L1Summary + | L1SkippedUpToDate + deriving stock (Generic, Eq, Show) + +-- | What the pass actually did with an atom, as opposed to what the +-- consolidator asked for. A merge whose targets all turned out to be missing +-- degrades to a store; a merge whose only target is the atom's own prior copy +-- degrades to a skip. The audit row records this, not the LLM's claim. +data AppliedAction + = ActionStored + | ActionUpdated + | ActionMerged + | ActionSkipped + deriving stock (Generic, Eq, Show) + +data AppliedDecision = AppliedDecision + { appliedAction :: !AppliedAction, + winnerId :: !(Maybe MemoryId), + appliedTargets :: ![MemoryId], + appliedNote :: !(Maybe Text) + } + +data WatermarkRow = WatermarkRow + { sessionId :: !Text, + lastTurnIndex :: !Int32, + distilledAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data AuditRow = AuditRow + { decisionId :: !Text, + sessionId :: !Text, + namespace :: !Text, + scopeKind :: !(Maybe Text), + scopeRef :: !(Maybe Text), + candidateContent :: !Text, + decision :: !Text, + targetIds :: ![Text], + resultMemoryId :: !(Maybe Text), + rationale :: !(Maybe Text), + decidedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +-- | Run one L1 distillation pass. Under 'RespectWatermark' a session whose +-- highest turn index is already covered by a previous fully successful pass is +-- skipped before any LLM call, which is what makes keiro's at-least-once timer +-- re-fires cheap. The watermark advances only when the whole fold succeeds, so +-- a failed pass is retried in full. +distillSessionL1 :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + L1RunMode -> + DistillRuntime -> + FindMergeCandidates es -> + SessionId -> + Eff es (Either L1Error L1Outcome) +distillSessionL1 mode rt finder sid = do + sessionResult <- Session.getById sid + case sessionResult of + Left err -> pure (Left (L1SessionReadFailed err)) + Right Nothing -> pure (Left (L1SessionNotFound sid)) + Right (Just session) -> do + turnsResult <- Session.getTurns sid + case turnsResult of + Left err -> pure (Left (L1TurnReadFailed err)) + Right turns -> do + let maxTurnIndex = maximum (0 : fmap (.turnIndex) turns) + upToDate <- watermarkCovers mode sid maxTurnIndex + if upToDate + then pure (Right L1SkippedUpToDate) + else do + inputResult <- buildExtractInput sid session turns + case inputResult of + Left err -> pure (Left err) + Right input -> do + extractedResult <- liftIO (runExtraction rt input) + case extractedResult of + Left err -> pure (Left (L1ExtractionFailed (Text.pack (show err)))) + Right output -> do + foldResult <- + foldM + (stepAtom maxTurnIndex session) + (Right emptySummary {extracted = length output.atoms}) + output.atoms + case foldResult of + Left err -> pure (Left err) + Right summary -> do + writeWatermark sid maxTurnIndex + pure (Right (L1Distilled summary)) + where + stepAtom _ _ (Left err) _ = pure (Left err) + stepAtom maxTurnIndex session (Right summary) atom = + applyAtom rt finder sid session maxTurnIndex summary atom + +watermarkCovers :: + (Store :> es) => + L1RunMode -> + SessionId -> + Int -> + Eff es Bool +watermarkCovers IgnoreWatermark _ _ = pure False +watermarkCovers RespectWatermark sid maxTurnIndex = do + stored <- readWatermark sid + pure (maybe False (>= maxTurnIndex) stored) + +readWatermark :: + (Store :> es) => + SessionId -> + Eff es (Maybe Int) +readWatermark sid = + runTransaction $ + fmap fromIntegral <$> Tx.statement (idText sid) selectWatermarkStmt + +writeWatermark :: + (IOE :> es, Store :> es) => + SessionId -> + Int -> + Eff es () +writeWatermark sid maxTurnIndex = do + now <- liftIO getCurrentTime + runTransaction $ + Tx.statement + WatermarkRow + { sessionId = idText sid, + lastTurnIndex = fromIntegral maxTurnIndex, + distilledAt = now + } + upsertWatermarkStmt + +scopedScanCandidates :: + (IOE :> es, Store :> es) => + Int -> + FindMergeCandidates es +scopedScanCandidates limit = + FindMergeCandidates \scope _query -> + fmap (take (max 0 limit)) <$> Recall.getActiveByScope scope + +recallCandidates :: + (IOE :> es, Store :> es) => + EmbeddingModel -> + VectorCapability -> + Int -> + FindMergeCandidates es +recallCandidates model capability limit = + FindMergeCandidates \scope query -> do + hits <- + Recall.recall + model + capability + Recall.RecallRequest + { scope, + query, + strategy = Recall.Hybrid, + maxResults = max 0 limit + } + pure (Right (fmap (.memory) hits)) + +buildExtractInput :: + (IOE :> es, Store :> es) => + SessionId -> + SessionRow -> + [TurnRow] -> + Eff es (Either L1Error ExtractInput) +buildExtractInput sid session turns = do + memoryTextResult <- + if null turns + then fallbackMemoryText sid (sessionScope session) + else pure (Right (renderTurns turns)) + pure do + memoryText <- memoryTextResult + Right + ExtractInput + { focus = field session.focus, + scopeLabel = field (renderScope (sessionScope session)), + conversation = field memoryText + } + +fallbackMemoryText :: + (IOE :> es, Store :> es) => + SessionId -> + MemoryScope -> + Eff es (Either L1Error Text) +fallbackMemoryText sid scope = do + bySession <- Recall.getBySession sid + case bySession of + Left err -> pure (Left (L1MemoryReadFailed err)) + Right rows + | not (null rows) -> pure (Right (renderMemories rows)) + | otherwise -> do + byScope <- Recall.getActiveByScope scope + pure $ + case byScope of + Left err -> Left (L1MemoryReadFailed err) + Right scopeRows -> Right (renderMemories scopeRows) + +applyAtom :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + DistillRuntime -> + FindMergeCandidates es -> + SessionId -> + SessionRow -> + Int -> + L1Summary -> + ExtractedAtom -> + Eff es (Either L1Error L1Summary) +applyAtom rt finder sid session maxTurnIndex summary atom = do + candidatesResult <- finder.runFindMergeCandidates (sessionScope session) (unField atom.content) + case candidatesResult of + Left err -> pure (Left (L1MemoryReadFailed err)) + Right candidates -> do + decisionResult <- + liftIO $ + runConsolidation + rt + ConsolidateInput + { scopeLabel = field (renderScope (sessionScope session)), + candidate = atom, + existing = existingMemory <$> candidates + } + case decisionResult of + Left err -> pure (Left (L1ConsolidationFailed (Text.pack (show err)))) + Right decision -> do + appliedResult <- applyDecision sid session atom decision + case appliedResult of + Left err -> pure (Left err) + Right applied -> do + writeAudit sid session atom maxTurnIndex decision applied + pure (Right (addAppliedDecision summary applied)) + +-- | Apply a consolidation decision, writing nothing until the whole plan for +-- the atom is known to be executable. The winner id is a deterministic function +-- of the session and the atom content, so every write here is idempotent under +-- keiro's at-least-once timer contract. +applyDecision :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + SessionRow -> + ExtractedAtom -> + ConsolidationDecision -> + Eff es (Either L1Error AppliedDecision) +applyDecision sid session atom decision = + case decision.action of + SkipAtom -> pure (Right (appliedSkip Nothing)) + StoreAtom -> storeWinner Nothing Nothing + UpdateAtom -> mergeInto ActionUpdated + MergeAtom -> mergeInto ActionMerged + where + winner = l1AtomMemoryId sid (unField atom.content) + + appliedSkip note = + AppliedDecision + { appliedAction = ActionSkipped, + winnerId = Nothing, + appliedTargets = [], + appliedNote = note + } + + storeWinner note supersedes = + fmap + ( \mid -> + AppliedDecision + { appliedAction = ActionStored, + winnerId = Just mid, + appliedTargets = [], + appliedNote = note + } + ) + <$> recordAtom sid session atom decision winner supersedes + + mergeInto action = do + let requested = nub (parsedTargetIds decision) + nonSelf = filter (/= winner) requested + degradeNote + | null requested = "no usable merge targets supplied; stored the candidate" + | otherwise = "targets missing; degraded to store" + if not (null requested) && null nonSelf + then pure (Right (appliedSkip (Just selfTargetNote))) + else do + winnerRow <- Memory.getMemoryRowById winner + case winnerRow of + Left err -> pure (Left (L1MemoryReadFailed err)) + Right (Just row) + | row.status /= "active" -> + pure (Right (appliedSkip (Just (retiredWinnerNote row.status)))) + _ -> do + resolved <- resolveExistingTargets nonSelf + case resolved of + Left err -> pure (Left err) + Right [] -> storeWinner (Just degradeNote) Nothing + Right targets@(firstTarget : _) -> do + winnerResult <- recordAtom sid session atom decision winner (Just firstTarget) + case winnerResult of + Left err -> pure (Left err) + Right stored -> do + mergeResults <- traverse (\target -> requireMemoryWrite =<< Memory.merge target stored) targets + pure $ + case lefts mergeResults of + err : _ -> Left err + [] -> + Right + AppliedDecision + { appliedAction = action, + winnerId = Just stored, + appliedTargets = targets, + appliedNote = Nothing + } + +-- | Drop target ids that name no row in the read model. A hallucinated but +-- syntactically valid TypeID would otherwise fail @Memory.merge@ with +-- 'Memory.MemoryNotFound' /after/ the winner had already been recorded, +-- wedging the timer and leaking one memory per retry. +resolveExistingTargets :: + (IOE :> es, Store :> es) => + [MemoryId] -> + Eff es (Either L1Error [MemoryId]) +resolveExistingTargets = + foldM step (Right []) + where + step (Left err) _ = pure (Left err) + step (Right acc) mid = do + row <- Memory.getMemoryRowById mid + pure $ + case row of + Left err -> Left (L1MemoryReadFailed err) + Right Nothing -> Right acc + Right (Just _) -> Right (acc <> [mid]) + +selfTargetNote :: Text +selfTargetNote = + "merge target is the atom's own prior copy; already represented" + +retiredWinnerNote :: Text -> Text +retiredWinnerNote status = + "this atom was already distilled and is now " <> status <> "; already represented" + +recordAtom :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + SessionRow -> + ExtractedAtom -> + ConsolidationDecision -> + MemoryId -> + Maybe MemoryId -> + Eff es (Either L1Error MemoryId) +recordAtom sid session atom decision memoryId supersedes = do + now <- liftIO getCurrentTime + requireMemoryWrite + =<< Memory.record + RecordMemoryData + { memoryId, + agentId = session.agentId, + sessionId = Just sid, + scope = sessionScope session, + memoryType = atomMemoryType atom, + content = resolvedContent atom decision, + priority = unField atom.priority, + confidence = atomConfidence atom, + tags = Set.fromList ["distilled", "l1"], + supersedes, + recordedAt = now + } + +requireMemoryWrite :: (Applicative f) => Either Memory.MemoryWriteError MemoryId -> f (Either L1Error MemoryId) +requireMemoryWrite = \case + Left err -> pure (Left (L1MemoryWriteFailed err)) + Right mid -> pure (Right mid) + +-- | The audit key is a deterministic function of the session, the pass's +-- maximum turn index, and the atom content, so @ON CONFLICT (decision_id) DO +-- NOTHING@ collapses re-fires of the same pass into one row while a later pass +-- over new turns still writes its own. +writeAudit :: + (IOE :> es, Store :> es) => + SessionId -> + SessionRow -> + ExtractedAtom -> + Int -> + ConsolidationDecision -> + AppliedDecision -> + Eff es () +writeAudit sid session atom maxTurnIndex decision applied = do + now <- liftIO getCurrentTime + runTransaction $ + Tx.statement + AuditRow + { decisionId = l1AuditKey sid maxTurnIndex (unField atom.content), + sessionId = idText sid, + namespace = scopeNamespaceText (sessionScope session), + scopeKind = scopeKindText (sessionScope session), + scopeRef = scopeRefText (sessionScope session), + candidateContent = unField atom.content, + decision = appliedActionText applied.appliedAction, + targetIds = idText <$> applied.appliedTargets, + resultMemoryId = idText <$> applied.winnerId, + rationale = Just (auditRationale decision applied), + decidedAt = now + } + insertAuditStmt + +auditRationale :: ConsolidationDecision -> AppliedDecision -> Text +auditRationale decision applied = + unField decision.rationale <> maybe "" (\note -> " [" <> note <> "]") applied.appliedNote + +-- | UUIDv5 namespace for every deterministic L1 identity (atom memory ids and +-- consolidation audit keys). +l1AtomNamespace :: UUID +l1AtomNamespace = + fromMaybe UUID.nil $ + UUID.fromString "6b696f6b-752d-6c31-8000-61746f6d6964" + +l1Uuid5 :: Text -> UUID +l1Uuid5 = + UUIDv5.generateNamed l1AtomNamespace . BS.unpack . TE.encodeUtf8 + +-- | The memory id an extracted atom always maps to. Keyed on the /candidate/ +-- content rather than the consolidator's rewritten @resultContent@, so the id +-- stays stable across retries whose rewrites differ. +l1AtomMemoryId :: SessionId -> Text -> MemoryId +l1AtomMemoryId sid content = + KindID.decorateKindID (l1Uuid5 (idText sid <> ":" <> content)) + +l1AuditKey :: SessionId -> Int -> Text -> Text +l1AuditKey sid maxTurnIndex content = + "kioku_consolidation_decision:" + <> UUID.toText + (l1Uuid5 ("audit:" <> idText sid <> ":" <> Text.pack (show maxTurnIndex) <> ":" <> content)) + +existingMemory :: MemoryRecord -> ExistingMemory +existingMemory row = + ExistingMemory + { memoryId = field row.memoryId, + memoryType = field row.memoryType, + content = field row.content, + priority = field row.priority, + confidence = field row.confidence + } + +parsedTargetIds :: ConsolidationDecision -> [MemoryId] +parsedTargetIds decision = + mapMaybe parseTargetId decision.targetMemoryIds + +parseTargetId :: Text -> Maybe MemoryId +parseTargetId = + either (const Nothing) Just . parseIdLenient + +resolvedContent :: ExtractedAtom -> ConsolidationDecision -> Text +resolvedContent atom decision = + fromMaybe (unField atom.content) decision.resultContent + +atomMemoryType :: ExtractedAtom -> MemoryType +atomMemoryType atom = + fromMaybe MemoryFact $ + memoryTypeFromText (Text.toLower (unField atom.atomType)) + +atomConfidence :: ExtractedAtom -> Confidence +atomConfidence atom = + fromMaybe MediumConfidence $ + confidenceFromText (Text.toLower (unField atom.confidence)) + +sessionScope :: SessionRow -> MemoryScope +sessionScope session = + scopeFromColumns session.namespace session.scopeKind session.scopeRef + +renderScope :: MemoryScope -> Text +renderScope scope = + Text.intercalate "/" $ + scopeNamespaceText scope : catMaybes [scopeKindText scope, scopeRefText scope] + +renderTurns :: [TurnRow] -> Text +renderTurns = + Text.intercalate "\n" . fmap renderTurn + +renderTurn :: TurnRow -> Text +renderTurn turn = + Text.pack (show turn.turnIndex) + <> ". " + <> turn.role + <> ": " + <> turn.content + <> maybe "" ("\n tools: " <>) turn.toolSummary + +renderMemories :: [MemoryRecord] -> Text +renderMemories = + Text.intercalate "\n" . fmap renderMemory + +renderMemory :: MemoryRecord -> Text +renderMemory row = + "- " <> row.memoryType <> ": " <> row.content + +addAppliedDecision :: L1Summary -> AppliedDecision -> L1Summary +addAppliedDecision summary applied = + case applied.appliedAction of + ActionStored -> summary {stored = summary.stored + 1} + ActionUpdated -> summary {merged = summary.merged + 1} + ActionMerged -> summary {merged = summary.merged + 1} + ActionSkipped -> summary {skipped = summary.skipped + 1} + +appliedActionText :: AppliedAction -> Text +appliedActionText = \case + ActionStored -> "store" + ActionUpdated -> "update" + ActionMerged -> "merge" + ActionSkipped -> "skip" + +emptySummary :: L1Summary +emptySummary = L1Summary {extracted = 0, stored = 0, merged = 0, skipped = 0} + +encodeTargetIds :: [Text] -> Text +encodeTargetIds = + TE.decodeUtf8 . BL.toStrict . Aeson.encode + +selectWatermarkStmt :: Statement Text (Maybe Int32) +selectWatermarkStmt = + preparable + """ + SELECT last_turn_index + FROM kioku_l1_watermarks + WHERE session_id = $1 + """ + (E.param (E.nonNullable E.text)) + (D.rowMaybe (D.column (D.nonNullable D.int4))) + +-- | @GREATEST@ keeps the watermark monotonic: a slow pass over turns 1-3 that +-- lands after a fast pass over turns 1-5 must not rewind it. +upsertWatermarkStmt :: Statement WatermarkRow () +upsertWatermarkStmt = + preparable + """ + INSERT INTO kioku_l1_watermarks (session_id, last_turn_index, distilled_at) + VALUES ($1, $2, $3) + ON CONFLICT (session_id) DO UPDATE + SET last_turn_index = + GREATEST(kioku_l1_watermarks.last_turn_index, EXCLUDED.last_turn_index), + distilled_at = EXCLUDED.distilled_at + """ + watermarkRowEncoder + D.noResult + +watermarkRowEncoder :: E.Params WatermarkRow +watermarkRowEncoder = + ((\row -> row.sessionId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.lastTurnIndex) >$< E.param (E.nonNullable E.int4)) + <> ((\row -> row.distilledAt) >$< E.param (E.nonNullable E.timestamptz)) + +insertAuditStmt :: Statement AuditRow () +insertAuditStmt = + preparable + """ + INSERT INTO kioku_consolidation_decisions + (decision_id, session_id, namespace, scope_kind, scope_ref, candidate_content, + decision, target_ids, result_memory_id, rationale, decided_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8::jsonb, $9, $10, $11) + ON CONFLICT (decision_id) DO NOTHING + """ + auditRowEncoder + D.noResult + +auditRowEncoder :: E.Params AuditRow +auditRowEncoder = + ((\row -> row.decisionId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.sessionId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.namespace) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.scopeKind) >$< E.param (E.nullable E.text)) + <> ((\row -> row.scopeRef) >$< E.param (E.nullable E.text)) + <> ((\row -> row.candidateContent) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.decision) >$< E.param (E.nonNullable E.text)) + <> ((encodeTargetIds . \row -> row.targetIds) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.resultMemoryId) >$< E.param (E.nullable E.text)) + <> ((\row -> row.rationale) >$< E.param (E.nullable E.text)) + <> ((\row -> row.decidedAt) >$< E.param (E.nonNullable E.timestamptz)) diff --git a/src/Kioku/Distill/L2.hs b/src/Kioku/Distill/L2.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/L2.hs @@ -0,0 +1,519 @@ +{-# LANGUAGE DataKinds #-} + +module Kioku.Distill.L2 + ( L2Error (..), + SceneRow (..), + fireL2SceneTimer, + getScenesByScope, + l2SceneProcessManagerName, + l2SceneTimerId, + l2SceneTimerScheduleProjection, + mirrorSceneToCurrentWorkspace, + mirrorSceneToWorkspace, + regenerateScene, + sceneMirrorPath, + sceneRowId, + ) +where + +import Contravariant.Extras (contrazip3, contrazip4) +import Control.Exception (IOException, try) +import Crypto.Hash (Digest, SHA256) +import Crypto.Hash qualified as Hash +import Data.Aeson qualified as Aeson +import Data.ByteString qualified as BS +import Data.ByteString.Lazy qualified as BL +import Data.Foldable (for_) +import Data.Functor.Contravariant ((>$<)) +import Data.Maybe (catMaybes) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE +import Data.Text.IO qualified as TextIO +import Data.Time (NominalDiffTime, addUTCTime) +import Data.UUID (UUID) +import Data.UUID qualified as UUID +import Data.UUID.V5 qualified as UUIDv5 +import Effectful (Eff, IOE, (:>)) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Keiro.Projection (InlineProjection (..)) +import Keiro.ReadModel (ReadModelError) +import Keiro.Timer (TimerId (..), TimerRequest (..), TimerRow (..), scheduleTimerTx) +import Kioku.Api.Scope (MemoryScope, scopeFromColumns, scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Api.Types (MemoryRecord (..)) +import Kioku.Distill.L3 (scheduleL3PersonaTimerTx) +import Kioku.Distill.Runtime (DistillRuntime, distillWorkspaceRoot, runSceneDistillation) +import Kioku.Distill.Scene (SceneInput (..), SceneOutput (..)) +import Kioku.Distill.ScopeIdentity (escapeScopeComponent, scopeIdentity, scopeSlugFromColumns) +import Kioku.Distill.Timer.Outcome (FireOutcome (..), fireRetryDelay, timerMarkerEventId) +import Kioku.Id (MemoryId, idText) +import Kioku.Memory.Domain + ( MemoryArchivedData (..), + MemoryConfidenceUpdatedData (..), + MemoryEvent (..), + MemoryMergedData (..), + MemoryRecordedData (..), + MemorySupersededData (..), + ) +import Kioku.Prelude +import Kioku.Recall qualified as Recall +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Transaction (runTransaction) +import Kiroku.Store.Types (EventId (..), RecordedEvent (..)) +import Shikumi.Schema.Types (field, unField) +import System.Directory (createDirectoryIfMissing, doesFileExist, getCurrentDirectory, removeFile) +import System.FilePath (()) + +data L2Error + = L2MemoryReadFailed !ReadModelError + | L2SceneReadFailed + | L2SceneGenerationFailed !Text + deriving stock (Generic, Show) + +data SceneRow = SceneRow + { sceneId :: !Text, + namespace :: !Text, + scopeKind :: !(Maybe Text), + scopeRef :: !(Maybe Text), + sceneKey :: !Text, + title :: !Text, + bodyMd :: !Text, + atomIds :: ![Text], + sourceHash :: !Text, + createdAt :: !UTCTime, + updatedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +newtype SceneTimerPayload = SceneTimerPayload + { scope :: MemoryScope + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +l2SceneProcessManagerName :: Text +l2SceneProcessManagerName = "kioku-l2-scene" + +defaultSceneKey :: Text +defaultSceneKey = "default" + +sceneDebounceSeconds :: NominalDiffTime +sceneDebounceSeconds = 5 + +l2SceneTimerScheduleProjection :: InlineProjection MemoryEvent +l2SceneTimerScheduleProjection = + InlineProjection + { name = "kioku-l2-scene-timer-schedule", + apply = scheduleSceneTimersForEvent + } + +-- | Every event that changes which memories a scope's scene is built from, or what +-- those memories say, must schedule a regeneration. Forgetting is such a change: +-- without this, archived, superseded, and merged content survives in the scene row +-- and its plaintext mirror until some unrelated memory happens to be recorded in +-- the same scope. So is a confidence change, for the same reason. +scheduleSceneTimersForEvent :: MemoryEvent -> RecordedEvent -> Tx.Transaction () +scheduleSceneTimersForEvent event recorded = case event of + MemoryRecorded d -> + scheduleTimerTx $ + l2SceneTimerRequest + d.scope + (idText (d.memoryId :: MemoryId)) + (addUTCTime sceneDebounceSeconds d.recordedAt) + MemoryArchived d -> + scheduleScopedSceneTimerTx d.memoryId (kindSourceId d.memoryId "archived") d.archivedAt + MemorySuperseded d -> + scheduleScopedSceneTimerTx d.memoryId (kindSourceId d.memoryId "superseded") d.supersededAt + MemoryMerged d -> + scheduleScopedSceneTimerTx d.memoryId (kindSourceId d.memoryId "merged") d.mergedAt + -- Confidence is in the scene's source hash ('atomSource') and in the LLM prompt + -- ('renderAtom'), so changing it makes the scene stale exactly as forgetting does. + -- + -- The source id must carry the event id. Unlike the forget events, which are + -- terminal, confidence can change repeatedly on one memory — and keiro re-arms a + -- conflicting timer only while it is still @scheduled@ (see 'scheduleScopedSceneTimerTx'). + -- A fixed @\:confidence@ would therefore be silently dropped on every + -- change after the first, which would look fixed and be worse than the bug. The + -- event id is stable across replay, so it cannot double-schedule either. + MemoryConfidenceUpdated d -> + scheduleScopedSceneTimerTx + d.memoryId + (idText d.memoryId <> ":confidence:" <> eventIdText recorded.eventId) + d.updatedAt + -- Tags are in neither 'atomSource' nor 'renderAtom', so a tag change cannot alter + -- the scene. Scheduling here would spend an LLM call to rewrite a byte-identical row. + MemoryTagsUpdated _ -> pure () + +-- | Schedule a scene regeneration for a memory whose event does not carry its scope. +-- +-- The scope is read back from the read-model row inside this same transaction. The +-- row is guaranteed present: the aggregate only accepts these commands for an +-- @Active@ memory, which implies a committed @MemoryRecorded@ whose inline +-- projection upserted the row. +-- +-- The caller supplies the source id rather than reusing the record path's bare +-- memory id, because keiro's 'scheduleTimerTx' re-arms a conflicting timer only +-- while it is still @scheduled@ — reusing the record-time id would be silently +-- dropped once that timer has fired, which by then it almost always has. +scheduleScopedSceneTimerTx :: MemoryId -> Text -> UTCTime -> Tx.Transaction () +scheduleScopedSceneTimerTx memoryId sourceId occurredAt = do + scopeCols <- Tx.statement (idText memoryId) selectMemoryScopeColumnsStmt + for_ scopeCols \(ns, sk, sr) -> + scheduleTimerTx $ + l2SceneTimerRequest + (scopeFromColumns ns sk sr) + sourceId + (addUTCTime sceneDebounceSeconds occurredAt) + +-- | The forget events are terminal, so one timer per (memory, event kind) is both +-- unique and stable across replay. +kindSourceId :: MemoryId -> Text -> Text +kindSourceId memoryId kind = idText memoryId <> ":" <> kind + +eventIdText :: EventId -> Text +eventIdText (EventId uuid) = UUID.toText uuid + +l2SceneTimerRequest :: MemoryScope -> Text -> UTCTime -> TimerRequest +l2SceneTimerRequest scope sourceId fireAt = + TimerRequest + { timerId = l2SceneTimerId scope sourceId, + processManagerName = l2SceneProcessManagerName, + correlationId = scopeIdentity scope, + fireAt, + payload = Aeson.toJSON (SceneTimerPayload scope) + } + +l2SceneTimerId :: MemoryScope -> Text -> TimerId +l2SceneTimerId scope sourceId = + TimerId $ + UUIDv5.generateNamed + l2SceneTimerNamespace + (BS.unpack (TE.encodeUtf8 raw)) + where + raw = + l2SceneProcessManagerName + <> ":" + <> scopeIdentity scope + <> ":" + <> sourceId + +regenerateScene :: + (IOE :> es, Store :> es) => + DistillRuntime -> + MemoryScope -> + Eff es (Either L2Error (Maybe SceneRow)) +regenerateScene rt scope = do + memoryResult <- Recall.getActiveByScope scope + case memoryResult of + Left err -> pure (Left (L2MemoryReadFailed err)) + -- Every memory in this scope has been forgotten. Delete the scene outright + -- rather than leaving it or blanking it: a surviving row keeps feeding + -- persona regeneration and keeps stale atom_ids pointing at forgotten + -- memories, and a blank mirror file is just a confusing way to still be + -- there. No LLM runs on this path -- there is nothing left to summarize. + Right [] -> do + existing <- lookupScene scope defaultSceneKey + case existing of + Left err -> pure (Left err) + Right Nothing -> pure (Right Nothing) + Right (Just row) -> do + now <- liftIO getCurrentTime + runTransaction do + Tx.statement row.sceneId deleteSceneStmt + scheduleL3PersonaTimerTx scope now + liftIO (bestEffortRemoveSceneMirror rt row) + pure (Right Nothing) + Right atoms -> do + let sourceHash = sceneSourceHash atoms + sceneId = sceneRowId scope + existing <- lookupScene scope defaultSceneKey + case existing of + Left err -> pure (Left err) + Right (Just row) + | row.sourceHash == sourceHash -> do + liftIO (bestEffortMirrorScene rt row) + pure (Right (Just row)) + _ -> do + outputResult <- + liftIO $ + runSceneDistillation + rt + SceneInput + { scopeLabel = field (renderScope scope), + atoms = field (renderAtoms atoms) + } + case outputResult of + Left err -> pure (Left (L2SceneGenerationFailed (Text.pack (show err)))) + Right output -> do + now <- liftIO getCurrentTime + let row = + SceneRow + { sceneId, + namespace = scopeNamespaceText scope, + scopeKind = scopeKindText scope, + scopeRef = scopeRefText scope, + sceneKey = defaultSceneKey, + title = unField output.title, + bodyMd = unField output.bodyMd, + atomIds = (.memoryId) <$> atoms, + sourceHash, + createdAt = now, + updatedAt = now + } + runTransaction do + Tx.statement row upsertSceneStmt + scheduleL3PersonaTimerTx scope now + liftIO (bestEffortMirrorScene rt row) + pure (Right (Just row)) + +fireL2SceneTimer :: + (IOE :> es, Store :> es) => + DistillRuntime -> + TimerRow -> + Eff es FireOutcome +fireL2SceneTimer rt row + | row.processManagerName /= l2SceneProcessManagerName = + pure FireNotMine + | otherwise = + case Aeson.fromJSON @SceneTimerPayload row.payload of + -- A payload this handler cannot parse will not parse on the next attempt + -- either. It used to be marked fired, which quietly lost the scene. + Aeson.Error err -> + pure (FireFailedPermanently ("L2 scene timer payload is malformed: " <> Text.pack err)) + Aeson.Success payload -> do + result <- regenerateScene rt payload.scope + pure $ + case result of + Right _ -> FireCompleted (timerMarkerEventId row.timerId) + Left err -> FireRetryLater (fireRetryDelay row.attempts) (Text.pack (show err)) + +lookupScene :: + (Store :> es) => + MemoryScope -> + Text -> + Eff es (Either L2Error (Maybe SceneRow)) +lookupScene scope sceneKey = do + result <- + runTransaction $ + Tx.statement + (scopeNamespaceText scope, scopeKindText scope, scopeRefText scope, sceneKey) + selectSceneByScopeKeyStmt + pure (Right result) + +getScenesByScope :: + (Store :> es) => + MemoryScope -> + Eff es [SceneRow] +getScenesByScope scope = + runTransaction $ + Tx.statement + (scopeNamespaceText scope, scopeKindText scope, scopeRefText scope) + selectScenesByScopeStmt + +mirrorSceneToCurrentWorkspace :: SceneRow -> IO FilePath +mirrorSceneToCurrentWorkspace row = do + workspace <- getCurrentDirectory + mirrorSceneToWorkspace workspace row + +mirrorSceneToWorkspace :: FilePath -> SceneRow -> IO FilePath +mirrorSceneToWorkspace workspace row = do + let path = sceneMirrorPath workspace row + createDirectoryIfMissing True (workspace ".kioku" "scenes") + TextIO.writeFile path (renderSceneFile row) + pure path + +sceneMirrorPath :: FilePath -> SceneRow -> FilePath +sceneMirrorPath workspace row = + workspace ".kioku" "scenes" Text.unpack (sceneScopeSlug row <> ".md") + +bestEffortMirrorScene :: DistillRuntime -> SceneRow -> IO () +bestEffortMirrorScene rt row = do + let write = do + workspace <- distillWorkspaceRoot rt + mirrorSceneToWorkspace workspace row + _ <- try write :: IO (Either IOException FilePath) + pure () + +-- | Remove a scene's mirror file, best-effort in exactly the way writing it is. +-- The durable artifact is the database row, and it is already gone by the time +-- this runs; a failure to unlink the file must not fail the timer, and the next +-- regeneration in this workspace rewrites or removes it anyway. +bestEffortRemoveSceneMirror :: DistillRuntime -> SceneRow -> IO () +bestEffortRemoveSceneMirror rt row = do + let remove = do + workspace <- distillWorkspaceRoot rt + let path = sceneMirrorPath workspace row + exists <- doesFileExist path + when exists (removeFile path) + _ <- try remove :: IO (Either IOException ()) + pure () + +renderSceneFile :: SceneRow -> Text +renderSceneFile row = + "# " <> row.title <> "\n\n" <> row.bodyMd <> "\n" + +sceneScopeSlug :: SceneRow -> Text +sceneScopeSlug row = + scopeSlugFromColumns row.namespace row.scopeKind row.scopeRef + +-- | The persistent primary key of a scene row. Escaped, so two distinct scopes can never +-- derive the same id; the scene key is escaped too, purely to future-proof the format +-- (today's only key, @default@, is unchanged by escaping). +sceneRowId :: MemoryScope -> Text +sceneRowId scope = + "kioku_scene:" <> scopeIdentity scope <> ":" <> escapeScopeComponent defaultSceneKey + +sceneSourceHash :: [MemoryRecord] -> Text +sceneSourceHash atoms = + "v1:" <> Text.pack (show (Hash.hash (BL.toStrict (Aeson.encode (atomSource <$> atoms))) :: Digest SHA256)) + +atomSource :: MemoryRecord -> (Text, Text, Int, Text, UTCTime) +atomSource atom = + (atom.memoryId, atom.content, atom.priority, atom.confidence, atom.createdAt) + +renderAtoms :: [MemoryRecord] -> Text +renderAtoms = + Text.intercalate "\n" . fmap renderAtom + +renderAtom :: MemoryRecord -> Text +renderAtom atom = + "- " + <> atom.memoryId + <> " (" + <> atom.memoryType + <> ", " + <> atom.confidence + <> "): " + <> atom.content + +-- | A human-readable scope label for the LLM prompt. Deliberately *not* escaped and +-- deliberately not used for identity: a collision here is cosmetic. Identity comes from +-- 'scopeIdentity'. +renderScope :: MemoryScope -> Text +renderScope scope = + Text.intercalate "/" $ + scopeNamespaceText scope : catMaybes [scopeKindText scope, scopeRefText scope] + +l2SceneTimerNamespace :: UUID +l2SceneTimerNamespace = + fromMaybe UUID.nil $ + UUID.fromString "6b696f6b-752d-7132-8000-7363656e6573" + +encodeAtomIds :: [Text] -> Text +encodeAtomIds = + TE.decodeUtf8 . BL.toStrict . Aeson.encode + +decodeAtomIds :: Text -> [Text] +decodeAtomIds = + fromMaybe [] . Aeson.decode . BL.fromStrict . TE.encodeUtf8 + +sceneRowDecoder :: D.Row SceneRow +sceneRowDecoder = + SceneRow + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> (decodeAtomIds <$> D.column (D.nonNullable D.text)) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.timestamptz) + <*> D.column (D.nonNullable D.timestamptz) + +sceneRowEncoder :: E.Params SceneRow +sceneRowEncoder = + ((\row -> row.sceneId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.namespace) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.scopeKind) >$< E.param (E.nullable E.text)) + <> ((\row -> row.scopeRef) >$< E.param (E.nullable E.text)) + <> ((\row -> row.sceneKey) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.title) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.bodyMd) >$< E.param (E.nonNullable E.text)) + <> ((encodeAtomIds . \row -> row.atomIds) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.sourceHash) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.createdAt) >$< E.param (E.nonNullable E.timestamptz)) + <> ((\row -> row.updatedAt) >$< E.param (E.nonNullable E.timestamptz)) + +selectMemoryScopeColumnsStmt :: Statement Text (Maybe (Text, Maybe Text, Maybe Text)) +selectMemoryScopeColumnsStmt = + preparable + "SELECT namespace, scope_kind, scope_ref FROM kioku_memories WHERE memory_id = $1" + (E.param (E.nonNullable E.text)) + ( D.rowMaybe + ( (,,) + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + ) + ) + +selectSceneByScopeKeyStmt :: Statement (Text, Maybe Text, Maybe Text, Text) (Maybe SceneRow) +selectSceneByScopeKeyStmt = + preparable + """ + SELECT scene_id, namespace, scope_kind, scope_ref, scene_key, title, body_md, + atom_ids::text, source_hash, created_at, updated_at + FROM kioku_scenes + WHERE namespace = $1 + AND ((scope_kind = $2 AND scope_ref = $3) + OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) + AND scene_key = $4 + """ + ( contrazip4 + (E.param (E.nonNullable E.text)) + (E.param (E.nullable E.text)) + (E.param (E.nullable E.text)) + (E.param (E.nonNullable E.text)) + ) + (D.rowMaybe sceneRowDecoder) + +selectScenesByScopeStmt :: Statement (Text, Maybe Text, Maybe Text) [SceneRow] +selectScenesByScopeStmt = + preparable + """ + SELECT scene_id, namespace, scope_kind, scope_ref, scene_key, title, body_md, + atom_ids::text, source_hash, created_at, updated_at + FROM kioku_scenes + WHERE namespace = $1 + AND ((scope_kind = $2 AND scope_ref = $3) + OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) + ORDER BY scene_key ASC, updated_at DESC + """ + ( contrazip3 + (E.param (E.nonNullable E.text)) + (E.param (E.nullable E.text)) + (E.param (E.nullable E.text)) + ) + (D.rowList sceneRowDecoder) + +-- | Delete by the row's own primary key, which was written from @sceneRowId@. +-- Deriving the key here instead would re-implement the scope-identity format +-- that docs/plans/13-... owns changing. +deleteSceneStmt :: Statement Text () +deleteSceneStmt = + preparable + "DELETE FROM kioku_scenes WHERE scene_id = $1" + (E.param (E.nonNullable E.text)) + D.noResult + +upsertSceneStmt :: Statement SceneRow () +upsertSceneStmt = + preparable + """ + INSERT INTO kioku_scenes + (scene_id, namespace, scope_kind, scope_ref, scene_key, title, body_md, + atom_ids, source_hash, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8::jsonb, $9, $10, $11) + ON CONFLICT (scene_id) DO UPDATE SET + title = EXCLUDED.title, + body_md = EXCLUDED.body_md, + atom_ids = EXCLUDED.atom_ids, + source_hash = EXCLUDED.source_hash, + updated_at = EXCLUDED.updated_at + """ + sceneRowEncoder + D.noResult diff --git a/src/Kioku/Distill/L3.hs b/src/Kioku/Distill/L3.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/L3.hs @@ -0,0 +1,386 @@ +{-# LANGUAGE DataKinds #-} + +module Kioku.Distill.L3 + ( L3Error (..), + PersonaRow (..), + fireL3PersonaTimer, + getPersonaByScope, + l3PersonaProcessManagerName, + l3PersonaTimerId, + mirrorPersonaToCurrentWorkspace, + mirrorPersonaToWorkspace, + personaMirrorPath, + personaRowId, + regeneratePersona, + scheduleL3PersonaTimerTx, + ) +where + +import Contravariant.Extras (contrazip3) +import Control.Exception (IOException, try) +import Crypto.Hash (Digest, SHA256) +import Crypto.Hash qualified as Hash +import Data.Aeson qualified as Aeson +import Data.ByteString qualified as BS +import Data.ByteString.Lazy qualified as BL +import Data.Functor.Contravariant ((>$<)) +import Data.Int (Int32) +import Data.Maybe (catMaybes) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE +import Data.Text.IO qualified as TextIO +import Data.Time (NominalDiffTime, addUTCTime) +import Data.UUID (UUID) +import Data.UUID qualified as UUID +import Data.UUID.V5 qualified as UUIDv5 +import Effectful (Eff, IOE, (:>)) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Keiro.Timer (TimerId (..), TimerRequest (..), TimerRow (..), scheduleTimerTx) +import Kioku.Api.Scope (MemoryScope, scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Distill.Persona (PersonaInput (..), PersonaOutput (..)) +import Kioku.Distill.Runtime (DistillRuntime, distillWorkspaceRoot, runPersonaDistillation) +import Kioku.Distill.ScopeIdentity (scopeIdentity, scopeSlugFromColumns) +import Kioku.Distill.Timer.Outcome (FireOutcome (..), fireRetryDelay, timerMarkerEventId) +import Kioku.Prelude +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Transaction (runTransaction) +import Shikumi.Schema.Types (field, unField) +import System.Directory (createDirectoryIfMissing, doesFileExist, getCurrentDirectory, removeFile) +import System.FilePath (()) + +data L3Error + = L3SceneGenerationUnavailable + | L3PersonaGenerationFailed !Text + deriving stock (Generic, Show) + +data PersonaRow = PersonaRow + { personaId :: !Text, + namespace :: !Text, + scopeKind :: !(Maybe Text), + scopeRef :: !(Maybe Text), + bodyMd :: !Text, + sceneCount :: !Int, + sourceHash :: !Text, + createdAt :: !UTCTime, + updatedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data PersonaSceneRow = PersonaSceneRow + { sceneId :: !Text, + title :: !Text, + bodyMd :: !Text, + updatedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +newtype PersonaTimerPayload = PersonaTimerPayload + { scope :: MemoryScope + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +l3PersonaProcessManagerName :: Text +l3PersonaProcessManagerName = "kioku-l3-persona" + +personaDebounceSeconds :: NominalDiffTime +personaDebounceSeconds = 5 + +scheduleL3PersonaTimerTx :: MemoryScope -> UTCTime -> Tx.Transaction () +scheduleL3PersonaTimerTx scope now = + scheduleTimerTx $ + TimerRequest + { timerId = l3PersonaTimerId scope fireAt, + processManagerName = l3PersonaProcessManagerName, + correlationId = scopeIdentity scope, + fireAt, + payload = Aeson.toJSON (PersonaTimerPayload scope) + } + where + fireAt = addUTCTime personaDebounceSeconds now + +l3PersonaTimerId :: MemoryScope -> UTCTime -> TimerId +l3PersonaTimerId scope fireAt = + TimerId $ + UUIDv5.generateNamed + l3PersonaTimerNamespace + (BS.unpack (TE.encodeUtf8 raw)) + where + raw = + l3PersonaProcessManagerName + <> ":" + <> scopeIdentity scope + <> ":" + <> Text.pack (show fireAt) + +regeneratePersona :: + (IOE :> es, Store :> es) => + DistillRuntime -> + MemoryScope -> + Eff es (Either L3Error (Maybe PersonaRow)) +regeneratePersona rt scope = do + scenes <- getPersonaScenesByScope scope + case scenes of + -- Every scene in this scope is gone, so the persona distilled from them has + -- no source left. Delete it and its mirror, symmetrically with the scene + -- delete in "Kioku.Distill.L2", and without an LLM call: there is nothing + -- to summarize. The persona is the top of the pyramid, so nothing chains on. + [] -> do + existing <- getPersonaByScope scope + case existing of + Nothing -> pure (Right Nothing) + Just row -> do + runTransaction (Tx.statement row.personaId deletePersonaStmt) + liftIO (bestEffortRemovePersonaMirror rt row) + pure (Right Nothing) + _ -> do + let sourceHash = personaSourceHash scenes + personaId = personaRowId scope + existing <- getPersonaByScope scope + case existing of + Just row + | row.sourceHash == sourceHash -> do + liftIO (bestEffortMirrorPersona rt row) + pure (Right (Just row)) + _ -> do + outputResult <- + liftIO $ + runPersonaDistillation + rt + PersonaInput + { scopeLabel = field (renderScope scope), + scenes = field (renderScenes scenes) + } + case outputResult of + Left err -> pure (Left (L3PersonaGenerationFailed (Text.pack (show err)))) + Right output -> do + now <- liftIO getCurrentTime + let row = + PersonaRow + { personaId, + namespace = scopeNamespaceText scope, + scopeKind = scopeKindText scope, + scopeRef = scopeRefText scope, + bodyMd = unField output.bodyMd, + sceneCount = length scenes, + sourceHash, + createdAt = now, + updatedAt = now + } + runTransaction (Tx.statement row upsertPersonaStmt) + liftIO (bestEffortMirrorPersona rt row) + pure (Right (Just row)) + +fireL3PersonaTimer :: + (IOE :> es, Store :> es) => + DistillRuntime -> + TimerRow -> + Eff es FireOutcome +fireL3PersonaTimer rt row + | row.processManagerName /= l3PersonaProcessManagerName = + pure FireNotMine + | otherwise = + case Aeson.fromJSON @PersonaTimerPayload row.payload of + -- Unparseable now, unparseable on every retry: dead-letter rather than + -- mark it fired and lose the persona silently. + Aeson.Error err -> + pure (FireFailedPermanently ("L3 persona timer payload is malformed: " <> Text.pack err)) + Aeson.Success payload -> do + result <- regeneratePersona rt payload.scope + pure $ + case result of + Right _ -> FireCompleted (timerMarkerEventId row.timerId) + Left err -> FireRetryLater (fireRetryDelay row.attempts) (Text.pack (show err)) + +getPersonaByScope :: + (Store :> es) => + MemoryScope -> + Eff es (Maybe PersonaRow) +getPersonaByScope scope = + runTransaction $ + Tx.statement + (scopeNamespaceText scope, scopeKindText scope, scopeRefText scope) + selectPersonaByScopeStmt + +getPersonaScenesByScope :: + (Store :> es) => + MemoryScope -> + Eff es [PersonaSceneRow] +getPersonaScenesByScope scope = + runTransaction $ + Tx.statement + (scopeNamespaceText scope, scopeKindText scope, scopeRefText scope) + selectScenesForPersonaStmt + +mirrorPersonaToCurrentWorkspace :: PersonaRow -> IO FilePath +mirrorPersonaToCurrentWorkspace row = do + workspace <- getCurrentDirectory + mirrorPersonaToWorkspace workspace row + +mirrorPersonaToWorkspace :: FilePath -> PersonaRow -> IO FilePath +mirrorPersonaToWorkspace workspace row = do + let path = personaMirrorPath workspace row + createDirectoryIfMissing True (workspace ".kioku" "persona") + TextIO.writeFile path (row.bodyMd <> "\n") + pure path + +personaMirrorPath :: FilePath -> PersonaRow -> FilePath +personaMirrorPath workspace row = + workspace ".kioku" "persona" Text.unpack (personaScopeSlug row <> ".md") + +bestEffortMirrorPersona :: DistillRuntime -> PersonaRow -> IO () +bestEffortMirrorPersona rt row = do + let write = do + workspace <- distillWorkspaceRoot rt + mirrorPersonaToWorkspace workspace row + _ <- try write :: IO (Either IOException FilePath) + pure () + +-- | Remove a persona's mirror file. Best-effort for the same reason writing it +-- is: the database row is the durable artifact and is already deleted. +bestEffortRemovePersonaMirror :: DistillRuntime -> PersonaRow -> IO () +bestEffortRemovePersonaMirror rt row = do + let remove = do + workspace <- distillWorkspaceRoot rt + let path = personaMirrorPath workspace row + exists <- doesFileExist path + when exists (removeFile path) + _ <- try remove :: IO (Either IOException ()) + pure () + +personaScopeSlug :: PersonaRow -> Text +personaScopeSlug row = + scopeSlugFromColumns row.namespace row.scopeKind row.scopeRef + +-- | The persistent primary key of a persona row. Escaped, so two distinct scopes can never +-- derive the same id. +personaRowId :: MemoryScope -> Text +personaRowId scope = + "kioku_persona:" <> scopeIdentity scope + +personaSourceHash :: [PersonaSceneRow] -> Text +personaSourceHash scenes = + "v1:" <> Text.pack (show (Hash.hash (BL.toStrict (Aeson.encode (sceneSource <$> scenes))) :: Digest SHA256)) + +sceneSource :: PersonaSceneRow -> (Text, Text, UTCTime) +sceneSource scene = + (scene.sceneId, scene.bodyMd, scene.updatedAt) + +renderScenes :: [PersonaSceneRow] -> Text +renderScenes = + Text.intercalate "\n\n" . fmap renderScene + +renderScene :: PersonaSceneRow -> Text +renderScene scene = + "# " <> scene.title <> "\n\n" <> scene.bodyMd + +-- | A human-readable scope label for the LLM prompt. Deliberately *not* escaped and +-- deliberately not used for identity: a collision here is cosmetic. Identity comes from +-- 'scopeIdentity'. +renderScope :: MemoryScope -> Text +renderScope scope = + Text.intercalate "/" $ + scopeNamespaceText scope : catMaybes [scopeKindText scope, scopeRefText scope] + +l3PersonaTimerNamespace :: UUID +l3PersonaTimerNamespace = + fromMaybe UUID.nil $ + UUID.fromString "6b696f6b-752d-7133-8000-706572736f6e" + +personaRowDecoder :: D.Row PersonaRow +personaRowDecoder = + PersonaRow + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nonNullable D.text) + <*> (fromIntegral @Int32 @Int <$> D.column (D.nonNullable D.int4)) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.timestamptz) + <*> D.column (D.nonNullable D.timestamptz) + +personaSceneRowDecoder :: D.Row PersonaSceneRow +personaSceneRowDecoder = + PersonaSceneRow + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.timestamptz) + +personaRowEncoder :: E.Params PersonaRow +personaRowEncoder = + ((\row -> row.personaId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.namespace) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.scopeKind) >$< E.param (E.nullable E.text)) + <> ((\row -> row.scopeRef) >$< E.param (E.nullable E.text)) + <> ((\row -> row.bodyMd) >$< E.param (E.nonNullable E.text)) + <> ((fromIntegral @Int @Int32 . \row -> row.sceneCount) >$< E.param (E.nonNullable E.int4)) + <> ((\row -> row.sourceHash) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.createdAt) >$< E.param (E.nonNullable E.timestamptz)) + <> ((\row -> row.updatedAt) >$< E.param (E.nonNullable E.timestamptz)) + +selectPersonaByScopeStmt :: Statement (Text, Maybe Text, Maybe Text) (Maybe PersonaRow) +selectPersonaByScopeStmt = + preparable + """ + SELECT persona_id, namespace, scope_kind, scope_ref, body_md, scene_count, + source_hash, created_at, updated_at + FROM kioku_personas + WHERE namespace = $1 + AND ((scope_kind = $2 AND scope_ref = $3) + OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) + """ + ( contrazip3 + (E.param (E.nonNullable E.text)) + (E.param (E.nullable E.text)) + (E.param (E.nullable E.text)) + ) + (D.rowMaybe personaRowDecoder) + +selectScenesForPersonaStmt :: Statement (Text, Maybe Text, Maybe Text) [PersonaSceneRow] +selectScenesForPersonaStmt = + preparable + """ + SELECT scene_id, title, body_md, updated_at + FROM kioku_scenes + WHERE namespace = $1 + AND ((scope_kind = $2 AND scope_ref = $3) + OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) + ORDER BY scene_key ASC, updated_at DESC + """ + ( contrazip3 + (E.param (E.nonNullable E.text)) + (E.param (E.nullable E.text)) + (E.param (E.nullable E.text)) + ) + (D.rowList personaSceneRowDecoder) + +-- | Delete by the row's own primary key, for the same reason 'deleteSceneStmt' +-- does: the scope-identity string format is not re-implemented here. +deletePersonaStmt :: Statement Text () +deletePersonaStmt = + preparable + "DELETE FROM kioku_personas WHERE persona_id = $1" + (E.param (E.nonNullable E.text)) + D.noResult + +upsertPersonaStmt :: Statement PersonaRow () +upsertPersonaStmt = + preparable + """ + INSERT INTO kioku_personas + (persona_id, namespace, scope_kind, scope_ref, body_md, scene_count, + source_hash, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) + ON CONFLICT (persona_id) DO UPDATE SET + body_md = EXCLUDED.body_md, + scene_count = EXCLUDED.scene_count, + source_hash = EXCLUDED.source_hash, + updated_at = EXCLUDED.updated_at + """ + personaRowEncoder + D.noResult diff --git a/src/Kioku/Distill/Persona.hs b/src/Kioku/Distill/Persona.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/Persona.hs @@ -0,0 +1,43 @@ +{-# LANGUAGE DataKinds #-} + +-- | Pure shikumi program for rendering L3 persona profiles from L2 scenes. +module Kioku.Distill.Persona + ( PersonaInput (..), + PersonaOutput (..), + personaSignature, + personaProgram, + ) +where + +import Kioku.Prelude +import Shikumi.Adapter (ToPrompt) +import Shikumi.Module (predict) +import Shikumi.Program (Program) +import Shikumi.Schema (FromModel, ToSchema, Validatable) +import Shikumi.Schema.Types (Field) +import Shikumi.Signature (Signature, mkSignature) + +data PersonaInput = PersonaInput + { scopeLabel :: Field "human label of the scope" Text, + scenes :: Field "the scene blocks for this scope, newline-joined" Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +newtype PersonaOutput = PersonaOutput + { bodyMd :: Field "a single distilled persona/profile markdown for this scope" Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt, Validatable) + +personaSignature :: Signature PersonaInput PersonaOutput +personaSignature = + mkSignature + "Distill all scene blocks for one scope into a single concise persona or \ + \profile. Describe who the agent is working with in this scope, their \ + \stable preferences, constraints, project facts, and durable patterns \ + \learned across scenes. Produce one markdown document. Preserve only details \ + \grounded in the scene text; do not invent biographical facts or private data." + +personaProgram :: Program PersonaInput PersonaOutput +personaProgram = predict personaSignature diff --git a/src/Kioku/Distill/Runtime.hs b/src/Kioku/Distill/Runtime.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/Runtime.hs @@ -0,0 +1,123 @@ +{-# LANGUAGE DataKinds #-} + +module Kioku.Distill.Runtime + ( DistillRuntime (..), + RuntimeSmokeInput (..), + RuntimeSmokeOutput (..), + distillWorkspaceRoot, + newDistillRuntime, + runConsolidation, + runDistillProgram, + runExtraction, + runPersonaDistillation, + runSceneDistillation, + runtimeSmokeProgram, + ) +where + +import Baikai.Model (Model) +import Baikai.Models.Generated qualified as Models +import Baikai.Provider.Claude.Api qualified as ClaudeApi +import Baikai.Provider.Registry (globalProviderRegistry) +import Effectful (runEff) +import Effectful.Concurrent (runConcurrent) +import Effectful.Error.Static (runErrorNoCallStack) +import Kioku.Distill.Consolidate (ConsolidateInput, ConsolidationDecision, consolidateProgram) +import Kioku.Distill.Extract (ExtractInput, ExtractOutput, extractProgram) +import Kioku.Distill.Persona (PersonaInput, PersonaOutput, personaProgram) +import Kioku.Distill.Scene (SceneInput, SceneOutput, sceneProgram) +import Kioku.Prelude +import Shikumi.Adapter (ToPrompt) +import Shikumi.Error (ShikumiError) +import Shikumi.LLM (LLMConfig, defaultLLMConfig, runLLMResilient) +import Shikumi.Module (predict) +import Shikumi.Program (Program, runProgram) +import Shikumi.Routing (routeLLM, runRouting) +import Shikumi.Schema (FromModel, ToSchema, Validatable) +import Shikumi.Schema.Types (Field) +import Shikumi.Signature (mkSignature) +import System.Directory (getCurrentDirectory) + +data DistillRuntime = DistillRuntime + { config :: !LLMConfig, + defaultModel :: !Model, + -- | Where the plaintext scene and persona mirrors are written and removed. + -- 'Nothing' means the process's working directory, which is what the CLI + -- wants: an agent's workspace is wherever it was invoked. Tests pin it to a + -- temp directory instead, because tasty runs cases concurrently and a + -- process-wide @chdir@ would race between them. + workspaceRoot :: !(Maybe FilePath), + runExtract :: !(ExtractInput -> IO (Either ShikumiError ExtractOutput)), + runConsolidate :: !(ConsolidateInput -> IO (Either ShikumiError ConsolidationDecision)), + runScene :: !(SceneInput -> IO (Either ShikumiError SceneOutput)), + runPersona :: !(PersonaInput -> IO (Either ShikumiError PersonaOutput)) + } + +distillWorkspaceRoot :: DistillRuntime -> IO FilePath +distillWorkspaceRoot rt = + maybe getCurrentDirectory pure rt.workspaceRoot + +newtype RuntimeSmokeInput = RuntimeSmokeInput + { prompt :: Field "short input text to echo" Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +newtype RuntimeSmokeOutput = RuntimeSmokeOutput + { answer :: Field "the input text repeated back" Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt, Validatable) + +newDistillRuntime :: IO DistillRuntime +newDistillRuntime = do + ClaudeApi.register + let config = defaultLLMConfig globalProviderRegistry + defaultModel = Models.anthropic_claude_haiku_4_5 + liveRun = runLiveDistillProgram config defaultModel + pure + DistillRuntime + { config, + defaultModel, + workspaceRoot = Nothing, + runExtract = liveRun extractProgram, + runConsolidate = liveRun consolidateProgram, + runScene = liveRun sceneProgram, + runPersona = liveRun personaProgram + } + +runDistillProgram :: DistillRuntime -> Program i o -> i -> IO (Either ShikumiError o) +runDistillProgram rt prog input = + runLiveDistillProgram rt.config rt.defaultModel prog input + +runExtraction :: DistillRuntime -> ExtractInput -> IO (Either ShikumiError ExtractOutput) +runExtraction rt = + rt.runExtract + +runConsolidation :: DistillRuntime -> ConsolidateInput -> IO (Either ShikumiError ConsolidationDecision) +runConsolidation rt = + rt.runConsolidate + +runSceneDistillation :: DistillRuntime -> SceneInput -> IO (Either ShikumiError SceneOutput) +runSceneDistillation rt = + rt.runScene + +runPersonaDistillation :: DistillRuntime -> PersonaInput -> IO (Either ShikumiError PersonaOutput) +runPersonaDistillation rt = + rt.runPersona + +runLiveDistillProgram :: LLMConfig -> Model -> Program i o -> i -> IO (Either ShikumiError o) +runLiveDistillProgram config model prog input = + runEff + . runErrorNoCallStack @ShikumiError + . runConcurrent + . runRouting model + . runLLMResilient config + . routeLLM + $ runProgram prog input + +runtimeSmokeProgram :: Program RuntimeSmokeInput RuntimeSmokeOutput +runtimeSmokeProgram = + predict $ + mkSignature + "Return the provided input text exactly once in the answer field." diff --git a/src/Kioku/Distill/Scene.hs b/src/Kioku/Distill/Scene.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/Scene.hs @@ -0,0 +1,44 @@ +{-# LANGUAGE DataKinds #-} + +-- | Pure shikumi program for rendering L2 scene blocks from L1 memory atoms. +module Kioku.Distill.Scene + ( SceneInput (..), + SceneOutput (..), + sceneSignature, + sceneProgram, + ) +where + +import Kioku.Prelude +import Shikumi.Adapter (ToPrompt) +import Shikumi.Module (predict) +import Shikumi.Program (Program) +import Shikumi.Schema (FromModel, ToSchema, Validatable) +import Shikumi.Schema.Types (Field) +import Shikumi.Signature (Signature, mkSignature) + +data SceneInput = SceneInput + { scopeLabel :: Field "human label of the scope" Text, + atoms :: Field "the active memory atoms in this scope, newline-joined" Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt) + +data SceneOutput = SceneOutput + { title :: Field "a short scene title, e.g. 'Testing & CI practices'" Text, + bodyMd :: Field "a markdown scene block summarizing the atoms as a narrative" Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToSchema, FromModel, ToPrompt, Validatable) + +sceneSignature :: Signature SceneInput SceneOutput +sceneSignature = + mkSignature + "Summarize a set of related agent memory atoms for one scope into a single \ + \readable markdown scene block. The title should name the dominant topic or \ + \workflow. The body should synthesize the atoms into a short narrative with \ + \a few bullet points or brief paragraphs. Preserve concrete preferences, \ + \constraints, and project facts; do not invent details not present in the atoms." + +sceneProgram :: Program SceneInput SceneOutput +sceneProgram = predict sceneSignature diff --git a/src/Kioku/Distill/ScopeIdentity.hs b/src/Kioku/Distill/ScopeIdentity.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/ScopeIdentity.hs @@ -0,0 +1,92 @@ +-- | Collision-free identity strings for a 'MemoryScope'. +-- +-- Distillation derives persistent row ids, timer ids, and mirror filenames from a scope's +-- namespace, kind, and ref. Those three are unconstrained host-supplied text, and the old +-- derivation simply joined them with @\/@: +-- +-- > ScopeGlobal (Namespace "a/b/c") -> "a/b/c" +-- > ScopeEntity (Namespace "a") (ScopeKind "b") "c" -> "a/b/c" +-- +-- Two different scopes, one identity — so both wrote to the same @kioku_scenes@ / +-- @kioku_personas@ row, and because the upserts do not update the scope columns on conflict, +-- the second scope's content landed on a row still attributed to the first. Cross-scope data +-- bleed, silently. +-- +-- The fix is to percent-escape each component before joining, which makes the encoding +-- injective per component and so the join unambiguous. Escaping @%@ first is what keeps it +-- injective: without that, an input already containing @%2F@ would decode ambiguously. +-- +-- Components containing none of @%@, @\/@ or @:@ — which is every well-formed scope in the +-- hosts and the docs (@rei:intention:intention_abc@, @mori:repo:web@, …) — encode to +-- themselves, so existing ids stay byte-identical and no mass migration is needed. +module Kioku.Distill.ScopeIdentity + ( escapeScopeComponent, + scopeIdentity, + scopeIdentityFromColumns, + scopeSlugFromColumns, + ) +where + +import Crypto.Hash (Digest, SHA256) +import Crypto.Hash qualified as Hash +import Data.Maybe (catMaybes) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE +import Kioku.Api.Scope (MemoryScope, scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Prelude + +-- | Percent-escape the three characters that would otherwise make the joined identity +-- ambiguous. @%@ must be escaped first, or the encoding is not injective. +escapeScopeComponent :: Text -> Text +escapeScopeComponent = + Text.replace ":" "%3A" + . Text.replace "/" "%2F" + . Text.replace "%" "%25" + +scopeIdentity :: MemoryScope -> Text +scopeIdentity scope = + scopeIdentityFromColumns + (scopeNamespaceText scope) + (scopeKindText scope) + (scopeRefText scope) + +-- | The row-column flavour. Kind and ref are both present or both absent — the +-- @*_scope_pair_check@ constraints added by the schema-hardening migration guarantee it — but +-- a half-populated row is treated as global here, matching 'Kioku.Api.Scope.scopeFromColumns'. +scopeIdentityFromColumns :: Text -> Maybe Text -> Maybe Text -> Text +scopeIdentityFromColumns namespace scopeKind scopeRef = + case (scopeKind, scopeRef) of + (Just kind, Just ref) -> + Text.intercalate "/" (escapeScopeComponent <$> [namespace, kind, ref]) + _ -> + escapeScopeComponent namespace + +-- | A workspace mirror filename: a human-readable prefix plus a hash of the true identity. +-- +-- The readable part alone cannot be collision-free — the sanitiser maps every unsafe +-- character to @-@, so @Namespace "a-b"@ and @Namespace "a" / ScopeKind "b"@ both render +-- @a-b@ — and no "escape only when ambiguous" rule can detect that locally. The suffix is +-- what actually separates them; the prefix is only so a human can tell the files apart. +scopeSlugFromColumns :: Text -> Maybe Text -> Maybe Text -> Text +scopeSlugFromColumns namespace scopeKind scopeRef = + sanitizeSlug readable <> "-" <> identityDigest + where + readable = + Text.intercalate "-" (namespace : catMaybes [scopeKind, scopeRef]) + + identityDigest = + Text.take 10 . Text.pack . show $ + (Hash.hash (TE.encodeUtf8 (scopeIdentityFromColumns namespace scopeKind scopeRef)) :: Digest SHA256) + +sanitizeSlug :: Text -> Text +sanitizeSlug = + Text.map \ch -> + if isSafeSlugChar ch then ch else '-' + +isSafeSlugChar :: Char -> Bool +isSafeSlugChar ch = + (ch >= 'a' && ch <= 'z') + || (ch >= 'A' && ch <= 'Z') + || (ch >= '0' && ch <= '9') + || ch == '-' + || ch == '_' diff --git a/src/Kioku/Distill/Timer.hs b/src/Kioku/Distill/Timer.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/Timer.hs @@ -0,0 +1,124 @@ +{-# LANGUAGE DataKinds #-} + +module Kioku.Distill.Timer + ( idleFlushSeconds, + l1ExtractProcessManagerName, + l1FinalTimerId, + l1IdleTimerId, + l1RampTimerId, + l1TimerScheduleProjection, + ) +where + +import Data.Aeson qualified as Aeson +import Data.ByteString qualified as BS +import Data.Foldable (traverse_) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE +import Data.Time (NominalDiffTime, addUTCTime) +import Data.UUID (UUID) +import Data.UUID qualified as UUID +import Data.UUID.V5 qualified as UUIDv5 +import Keiro.Projection (InlineProjection (..)) +import Keiro.Timer (TimerId (..), TimerRequest (..), scheduleTimerTx) +import Kioku.Id (SessionId, idText) +import Kioku.Prelude +import Kioku.Session.Domain + ( InteractiveSessionRecordedData (..), + SessionCompletedData (..), + SessionEvent (..), + SessionFailedData (..), + SessionStartedData (..), + TurnRecordedData (..), + ) + +idleFlushSeconds :: NominalDiffTime +idleFlushSeconds = 30 * 60 + +l1ExtractProcessManagerName :: Text +l1ExtractProcessManagerName = "kioku-l1-extract" + +l1TimerScheduleProjection :: InlineProjection SessionEvent +l1TimerScheduleProjection = + InlineProjection + { name = "kioku-l1-timer-schedule", + apply = \event _recorded -> traverse_ scheduleTimerTx (timerRequestsForEvent event) + } + +-- | Every timer id is a UUIDv5 over stable event data only, never over +-- @fireAt@, so re-projecting a session's events can never double-schedule. +-- +-- The idle id in particular is one per session: keiro's @scheduleTimerTx@ +-- upserts on @timer_id@ and re-arms the row (moving @fire_at@ forward) only +-- while it is still @scheduled@, so every recorded turn pushes the same single +-- row later instead of inserting another. That is the debounce — a 50-turn +-- session holds one idle timer, not fifty. +l1TimerIdFor :: Text -> TimerId +l1TimerIdFor suffix = + TimerId $ + UUIDv5.generateNamed + l1TimerNamespace + (BS.unpack (TE.encodeUtf8 (l1ExtractProcessManagerName <> ":" <> suffix))) + +l1IdleTimerId :: SessionId -> TimerId +l1IdleTimerId sid = l1TimerIdFor (idText sid <> ":idle") + +l1RampTimerId :: SessionId -> Int -> TimerId +l1RampTimerId sid turnIndex = + l1TimerIdFor (idText sid <> ":ramp:" <> Text.pack (show turnIndex)) + +l1FinalTimerId :: SessionId -> TimerId +l1FinalTimerId sid = l1TimerIdFor (idText sid <> ":final") + +timerRequestsForEvent :: SessionEvent -> [TimerRequest] +timerRequestsForEvent = \case + SessionStarted d -> + [idleRequest d.sessionId (addUTCTime idleFlushSeconds d.startedAt) (Just 0)] + InteractiveSessionRecorded d -> + [idleRequest d.sessionId (addUTCTime idleFlushSeconds d.startedAt) (Just 0)] + SessionCompleted d -> + [finalRequest d.sessionId d.completedAt] + SessionFailed d -> + [finalRequest d.sessionId (failedAt d)] + SessionAwaiting _ -> [] + SessionResumed _ -> [] + TurnRecorded d -> + let idle = idleRequest d.sessionId (addUTCTime idleFlushSeconds d.recordedAt) (Just d.turnIndex) + ramp = rampRequest d.sessionId d.turnIndex d.recordedAt + in if isRampTurn d.turnIndex then [ramp, idle] else [idle] + +idleRequest :: SessionId -> UTCTime -> Maybe Int -> TimerRequest +idleRequest sid = l1TimerRequest (l1IdleTimerId sid) sid "idle" + +rampRequest :: SessionId -> Int -> UTCTime -> TimerRequest +rampRequest sid turnIndex fireAt = + l1TimerRequest (l1RampTimerId sid turnIndex) sid "ramp" fireAt (Just turnIndex) + +finalRequest :: SessionId -> UTCTime -> TimerRequest +finalRequest sid fireAt = + l1TimerRequest (l1FinalTimerId sid) sid "final" fireAt Nothing + +-- | @turnCount@ in the payload is diagnostic only; the real freshness check is +-- the @kioku_l1_watermarks@ read inside the pass. +l1TimerRequest :: TimerId -> SessionId -> Text -> UTCTime -> Maybe Int -> TimerRequest +l1TimerRequest timerId sid kind fireAt turnCount = + TimerRequest + { timerId, + processManagerName = l1ExtractProcessManagerName, + correlationId = idText sid, + fireAt, + payload = + Aeson.object + [ "kind" Aeson..= kind, + "turnCount" Aeson..= turnCount + ] + } + +isRampTurn :: Int -> Bool +isRampTurn n = + n == 1 || n == 2 || n == 4 || n == 8 || n == 16 || (n > 16 && n `mod` 16 == 0) + +l1TimerNamespace :: UUID +l1TimerNamespace = + fromMaybe UUID.nil $ + UUID.fromString "6b696f6b-752d-7131-8000-6c3174696d72" diff --git a/src/Kioku/Distill/Timer/Outcome.hs b/src/Kioku/Distill/Timer/Outcome.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/Timer/Outcome.hs @@ -0,0 +1,67 @@ +-- | What a distillation timer fire decided. +-- +-- This module sits at the bottom of the distillation graph so that +-- "Kioku.Distill.L2", "Kioku.Distill.L3", and "Kioku.Distill.Timer.Worker" can +-- all import it (the worker imports L2 and L3, so the type cannot live there). +-- +-- It replaces a @Maybe EventId@ whose @Nothing@ meant three incompatible things +-- at once — a transient failure, a permanent failure, and "this timer is not +-- mine" — all of which keiro treated identically by leaving the row @firing@ +-- until its 300-second stale-claim requeue. A timer whose LLM call could never +-- succeed therefore retried forever, at full LLM cost, with nothing to show for +-- it. +module Kioku.Distill.Timer.Outcome + ( FireOutcome (..), + fireRetryDelay, + unknownTimerRetryDelay, + timerMarkerEventId, + ) +where + +import Data.Time (NominalDiffTime) +import Keiro.Timer (TimerId (..)) +import Kioku.Prelude +import Kiroku.Store.Types (EventId (..)) + +-- | The verdict of one fire attempt. +data FireOutcome + = -- | Done. Mark the timer fired with this event id. + FireCompleted !EventId + | -- | Transient failure: reschedule this many seconds out, logging the note. + -- Bounded by keiro's attempt ceiling, so a persistently failing timer ends + -- as a visible @dead@ row rather than cycling forever. + FireRetryLater !NominalDiffTime !Text + | -- | This can never succeed (a corrupt payload, an unparseable correlation + -- id). Dead-letter it with this reason instead of faking success. + FireFailedPermanently !Text + | -- | This timer's process manager is not mine; I did not touch the row. + FireNotMine + deriving stock (Generic, Eq, Show) + +-- | Backoff for a transient fire failure, by post-claim attempt count: +-- 30s, 60s, 120s, … doubling, capped at 900s. +-- +-- keiro increments @attempts@ at claim time, so the first failure is called with +-- @1@. Eight claims under this schedule span roughly an hour — long enough to +-- ride out a provider incident, short enough to stop burning LLM tokens on work +-- that is structurally broken. +fireRetryDelay :: Int -> NominalDiffTime +fireRetryDelay attempts = + min 900 (30 * (2 ^ max 0 (attempts - 1))) + +-- | Requeue delay for a timer no handler owns. +-- +-- keiro's 'Keiro.Timer.claimDueTimer' claims the earliest due timer regardless +-- of process-manager name, so an unknown-PM timer cannot be left unclaimed — it +-- must be put back. It is deliberately not dead-lettered on sight: during a +-- rolling deploy a newer kioku may have scheduled timers this binary does not +-- know about yet, and ten minutes gives the newer worker time to take them. The +-- attempt ceiling still guarantees a genuinely orphaned timer dies visibly. +unknownTimerRetryDelay :: NominalDiffTime +unknownTimerRetryDelay = 600 + +-- | Distillation timers append no dedicated domain event, so a fired timer is +-- marked with its own id as the event id. Keeping the convention in one place +-- means the three fire handlers cannot drift apart on it. +timerMarkerEventId :: TimerId -> EventId +timerMarkerEventId (TimerId uuid) = EventId uuid diff --git a/src/Kioku/Distill/Timer/Worker.hs b/src/Kioku/Distill/Timer/Worker.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Distill/Timer/Worker.hs @@ -0,0 +1,209 @@ +{-# LANGUAGE DataKinds #-} + +module Kioku.Distill.Timer.Worker + ( applyFireOutcome, + drainKiokuTimers, + fireKiokuTimer, + fireL1Timer, + kiokuTimerWorkerOptions, + runKiokuTimerWorkerOnce, + ) +where + +import Data.Text qualified as Text +import Data.Time (NominalDiffTime, addUTCTime) +import Effectful (Eff, IOE, (:>)) +import Effectful.Error.Static (Error) +import Keiro.Telemetry (KeiroMetrics) +import Keiro.Timer + ( TimerRequest (..), + TimerRow (..), + TimerWorkerOptions (..), + deadLetterTimer, + requeueStuckTimer, + runTimerWorkerWith, + scheduleTimerTx, + ) +import Kioku.Distill.L1 (FindMergeCandidates, L1Error (..), L1RunMode (..), distillSessionL1) +import Kioku.Distill.L2 (fireL2SceneTimer) +import Kioku.Distill.L3 (fireL3PersonaTimer) +import Kioku.Distill.Runtime (DistillRuntime) +import Kioku.Distill.Timer (l1ExtractProcessManagerName) +import Kioku.Distill.Timer.Outcome + ( FireOutcome (..), + fireRetryDelay, + timerMarkerEventId, + unknownTimerRetryDelay, + ) +import Kioku.Id (parseIdLenient) +import Kioku.Prelude +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Transaction (runTransaction) +import Kiroku.Store.Types (EventId (..)) +import System.IO qualified as IO + +-- | kioku's timer policy. +-- +-- Eight claims with 'fireRetryDelay''s backoff spans roughly an hour before a +-- timer is dead-lettered, which is the point of the ceiling: a structurally +-- failing distillation (a conversation past the model's context window, say) +-- must stop costing LLM tokens and start being visible instead. The 300-second +-- stale requeue is keiro's default and unchanged. +kiokuTimerWorkerOptions :: TimerWorkerOptions +kiokuTimerWorkerOptions = + TimerWorkerOptions {maxAttempts = Just 8, requeueStuckAfter = Just 300} + +fireL1Timer :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + DistillRuntime -> + FindMergeCandidates es -> + TimerRow -> + Eff es FireOutcome +fireL1Timer rt finder row + | row.processManagerName /= l1ExtractProcessManagerName = + pure FireNotMine + | otherwise = + case parseIdLenient row.correlationId of + -- A correlation id that is not a session id will never become one. + -- This used to be marked fired, which looked like success. + Left _err -> + pure + ( FireFailedPermanently + ("L1 timer correlation id is not a session id: " <> row.correlationId) + ) + Right sid -> do + result <- distillSessionL1 RespectWatermark rt finder sid + pure $ + case result of + -- Both a real pass and a watermark skip mean this timer is done. + Right _outcome -> FireCompleted (timerMarkerEventId row.timerId) + -- A session may legitimately be gone (deleted data); nothing to do. + Left (L1SessionNotFound _) -> FireCompleted (timerMarkerEventId row.timerId) + -- Everything else — a failed LLM extraction or consolidation, a + -- read-model error, a failed write — is worth another attempt, and + -- the attempt ceiling bounds how many. + Left err -> FireRetryLater (fireRetryDelay row.attempts) (Text.pack (show err)) + +-- | Offer the timer to each handler in turn. The handlers identify their own +-- work by process-manager name, so a 'FireNotMine' simply falls through to the +-- next one; a 'FireNotMine' from all three means no handler owns this timer, and +-- the runner decides what to do about that. +fireKiokuTimer :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + DistillRuntime -> + FindMergeCandidates es -> + TimerRow -> + Eff es FireOutcome +fireKiokuTimer rt finder row = do + l1Result <- fireL1Timer rt finder row + case l1Result of + FireNotMine -> do + l2Result <- fireL2SceneTimer rt row + case l2Result of + FireNotMine -> fireL3PersonaTimer rt row + outcome -> pure outcome + outcome -> pure outcome + +-- | Turn a fire verdict into keiro timer state. +-- +-- Returning 'Nothing' to keiro means "do not mark this row fired"; every such +-- branch has already moved the row itself, so the timer never sits in @firing@ +-- waiting on the 300-second stale requeue. +applyFireOutcome :: + (IOE :> es, Store :> es) => + TimerRow -> + FireOutcome -> + Eff es (Maybe EventId) +applyFireOutcome row = \case + FireCompleted eventId -> pure (Just eventId) + FireRetryLater delay note -> do + logTimer row ("retrying in " <> Text.pack (show delay) <> ": " <> note) + rescheduleClaimedTimer row delay + pure Nothing + FireFailedPermanently reason -> do + logTimer row ("dead-lettering: " <> reason) + void (deadLetterTimer row.timerId reason) + pure Nothing + FireNotMine -> do + logTimer row "no handler owns this process manager; requeueing" + rescheduleClaimedTimer row unknownTimerRetryDelay + pure Nothing + +-- | Push a claimed timer back out into the future. +-- +-- keiro has no single-call "reschedule a firing timer" at this pin, so this is +-- two public calls: 'requeueStuckTimer' moves the claimed @firing@ row back to +-- @scheduled@ (leaving @fire_at@ alone), and 'scheduleTimerTx' then re-arms it — +-- its upsert only updates rows that are @scheduled@, which is exactly what the +-- first call just guaranteed. @attempts@ is deliberately not reset, so the +-- ceiling counts total claims. +rescheduleClaimedTimer :: + (IOE :> es, Store :> es) => + TimerRow -> + NominalDiffTime -> + Eff es () +rescheduleClaimedTimer row delay = do + requeued <- requeueStuckTimer row.timerId + when requeued do + now <- liftIO getCurrentTime + runTransaction $ + scheduleTimerTx + TimerRequest + { timerId = row.timerId, + processManagerName = row.processManagerName, + correlationId = row.correlationId, + fireAt = addUTCTime delay now, + payload = row.payload + } + +runKiokuTimerWorkerOnce :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + Maybe KeiroMetrics -> + DistillRuntime -> + FindMergeCandidates es -> + UTCTime -> + Eff es (Maybe TimerRow) +runKiokuTimerWorkerOnce metrics rt finder now = + runTimerWorkerWith metrics kiokuTimerWorkerOptions now \row -> + fireKiokuTimer rt finder row >>= applyFireOutcome row + +-- | Claim and fire due timers until none remain, returning how many were +-- processed. +-- +-- The old loop slept between every timer, capping throughput at one timer per +-- poll interval. This cannot spin: every outcome other than 'FireCompleted' +-- either moves the row's @fire_at@ at least 30 seconds out or puts it in a +-- terminal state, so a timer processed in this pass is not claimable again +-- within it. +drainKiokuTimers :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + Maybe KeiroMetrics -> + DistillRuntime -> + FindMergeCandidates es -> + Eff es Int +drainKiokuTimers metrics rt finder = go 0 + where + go processed = do + now <- liftIO getCurrentTime + claimed <- runKiokuTimerWorkerOnce metrics rt finder now + case claimed of + Nothing -> pure processed + Just _ -> go (processed + 1) + +logTimer :: (IOE :> es) => TimerRow -> Text -> Eff es () +logTimer row msg = + liftIO + ( IO.hPutStrLn + IO.stderr + ( Text.unpack + ( "kioku-distill-timer[" + <> row.processManagerName + <> " attempt " + <> Text.pack (show row.attempts) + <> "]: " + <> msg + ) + ) + ) diff --git a/src/Kioku/Memory.hs b/src/Kioku/Memory.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Memory.hs @@ -0,0 +1,317 @@ +module Kioku.Memory + ( MemoryWriteError (..), + record, + supersede, + archive, + updateTags, + updateConfidence, + merge, + getMemoryRowById, + getActiveRowsInNamespace, + getActiveRowsByScope, + getRowsBySession, + getActiveRowsByType, + getSupersessionChain, + ) +where + +import Data.List (find) +import Effectful (Eff, IOE, (:>)) +import Effectful.Error.Static (Error) +import Keiro.Command (CommandError (..), defaultRunCommandOptions) +import Keiro.Projection (runCommandWithProjections) +import Keiro.ReadModel (ConsistencyMode (..), ReadModelError, runQueryWith) +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Api.Types (MemoryType, confidenceToText, memoryTypeToText) +import Kioku.Distill.L2 (l2SceneTimerScheduleProjection) +import Kioku.Id (MemoryId, SessionId, idText) +import Kioku.Memory.Domain +import Kioku.Memory.EventStream (memoryEventStream, memoryStream) +import Kioku.Memory.ReadModel + ( MemoriesByNamespaceQuery (..), + MemoriesByScopeQuery (..), + MemoriesBySessionQuery (..), + MemoriesByTypeQuery (..), + MemoryByIdQuery (..), + MemoryRow (..), + MemorySupersessionChainQuery (..), + memoriesByNamespaceRowsReadModel, + memoriesByScopeRowsReadModel, + memoriesBySessionRowsReadModel, + memoriesByTypeRowsReadModel, + memoryByIdReadModel, + memoryInlineProjection, + memorySupersessionChainReadModel, + ) +import Kioku.Prelude +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) + +data MemoryWriteError + = MemoryCommandRejected !CommandError + | MemoryReadFailed !ReadModelError + | MemoryNotFound + | MemoryNotActive + | MemoryConflict !Text + deriving stock (Generic, Show) + +record :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + RecordMemoryData -> + Eff es (Either MemoryWriteError MemoryId) +record cmdData = do + existing <- lookupMemory cmdData.memoryId + case existing of + Left err -> pure (Left (MemoryReadFailed err)) + Right (Just row) -> pure (idempotentOr "record" recordMismatch row cmdData.memoryId) + Right Nothing -> + runMemoryCommand cmdData.memoryId (RecordMemory cmdData) + >>= acceptRejectedIfMatches cmdData.memoryId (isNothing . recordMismatch) + where + recordMismatch = mismatchOf memoryRecordFields cmdData + +supersede :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SupersedeMemoryData -> + Eff es (Either MemoryWriteError MemoryId) +supersede cmdData = do + existing <- lookupMemory cmdData.memoryId + case existing of + Left err -> pure (Left (MemoryReadFailed err)) + Right Nothing -> pure (Left MemoryNotFound) + Right (Just row) + -- Already retired: only a supersession by the *same* winner is this request's own + -- echo. Superseding by a different winner is a conflict, not a duplicate. + | row.status /= "active" -> pure (idempotentOr "supersede" supersedeMismatch row cmdData.memoryId) + | otherwise -> + runMemoryCommand cmdData.memoryId (SupersedeMemory cmdData) + >>= acceptRejectedIfMatches cmdData.memoryId (isNothing . supersedeMismatch) + where + supersedeMismatch = mismatchOf memorySupersedeFields cmdData + +archive :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + ArchiveMemoryData -> + Eff es (Either MemoryWriteError MemoryId) +archive cmdData = do + existing <- lookupMemory cmdData.memoryId + case existing of + Left err -> pure (Left (MemoryReadFailed err)) + Right Nothing -> pure (Left MemoryNotFound) + Right (Just row) + | row.status /= "active" -> pure (idempotentOr "archive" archiveMismatch row cmdData.memoryId) + | otherwise -> + runMemoryCommand cmdData.memoryId (ArchiveMemory cmdData) + >>= acceptRejectedIfMatches cmdData.memoryId (isNothing . archiveMismatch) + where + archiveMismatch = mismatchOf memoryArchiveFields cmdData + +updateTags :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + UpdateMemoryTagsData -> + Eff es (Either MemoryWriteError MemoryId) +updateTags cmdData = do + existing <- lookupMemory cmdData.memoryId + case existing of + Left err -> pure (Left (MemoryReadFailed err)) + Right Nothing -> pure (Left MemoryNotFound) + Right (Just row) + | row.status /= "active" -> pure (Left MemoryNotActive) + | row.tags == cmdData.tags -> pure (Right cmdData.memoryId) + | otherwise -> runMemoryCommand cmdData.memoryId (UpdateMemoryTags cmdData) + +updateConfidence :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + UpdateMemoryConfidenceData -> + Eff es (Either MemoryWriteError MemoryId) +updateConfidence cmdData = do + existing <- lookupMemory cmdData.memoryId + case existing of + Left err -> pure (Left (MemoryReadFailed err)) + Right Nothing -> pure (Left MemoryNotFound) + Right (Just row) + | row.status /= "active" -> pure (Left MemoryNotActive) + | row.confidence == confidenceToText cmdData.confidence -> pure (Right cmdData.memoryId) + | otherwise -> runMemoryCommand cmdData.memoryId (UpdateMemoryConfidence cmdData) + +-- | Merge @loser@ into @winner@. +-- +-- Unlike the other writes, @mergedAt@ is generated here rather than supplied by the caller, +-- so a retry cannot re-deliver an identical timestamp. Idempotency therefore matches on the +-- merge target alone: merging into the same winner twice is a duplicate, merging into a +-- different one is a conflict. +merge :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + MemoryId -> + MemoryId -> + Eff es (Either MemoryWriteError MemoryId) +merge loser winner = do + existing <- lookupMemory loser + case existing of + Left err -> pure (Left (MemoryReadFailed err)) + Right Nothing -> pure (Left MemoryNotFound) + Right (Just row) + | row.status /= "active" -> pure (idempotentOr "merge" mergeMismatch row loser) + | otherwise -> do + now <- liftIO getCurrentTime + runMemoryCommand loser (MergeMemory (MergeMemoryData loser winner now)) + >>= acceptRejectedIfMatches loser (isNothing . mergeMismatch) + where + mergeMismatch = mismatchOf memoryMergeFields winner + +-- * Idempotent accepts + +-- | A named comparison between one request field and the memory row that already exists. +type FieldCheck cmd = (Text, cmd -> MemoryRow -> Bool) + +-- | The first request field that disagrees with the recorded row, if any. +-- +-- Call-time timestamps (@recordedAt@, @supersededAt@, @archivedAt@) are deliberately /not/ +-- compared. The id is the identity: a second write against the same id with the same +-- semantic payload is a retry, and retries re-read the clock. Distillation is the proof — +-- 'Kioku.Distill.L1.recordAtom' derives a deterministic memory id but passes +-- @recordedAt = now@, so comparing the timestamp would turn every idle-timer re-fire (the +-- exact regime L1's deterministic identity exists to survive) into a hard conflict. +-- +-- Everything that carries meaning — content, scope, type, priority, confidence, tags, +-- lineage, and the merge/supersession target — is compared, which is what the review +-- actually asked for: a reused id with different /content/ must not report success. +mismatchOf :: [FieldCheck cmd] -> cmd -> MemoryRow -> Maybe Text +mismatchOf checks cmd row = + fst <$> find (\(_, matches) -> not (matches cmd row)) checks + +-- | A duplicate request that matches what already happened succeeds; one that conflicts with +-- it gets a conflict error naming the field that differs. +idempotentOr :: + Text -> + (MemoryRow -> Maybe Text) -> + MemoryRow -> + MemoryId -> + Either MemoryWriteError MemoryId +idempotentOr operation mismatch row mid = + case mismatch row of + Nothing -> Right mid + Just field -> + Left (MemoryConflict (operation <> ": " <> field <> " differs from the recorded memory")) + +-- | Translate a losing concurrent-duplicate race into the success the winner got. See +-- 'Kioku.Session.acceptRejectedIfMatches' — same contract, memory side. +acceptRejectedIfMatches :: + (IOE :> es, Store :> es) => + MemoryId -> + (MemoryRow -> Bool) -> + Either MemoryWriteError MemoryId -> + Eff es (Either MemoryWriteError MemoryId) +acceptRejectedIfMatches mid matches = \case + Left err@(MemoryCommandRejected CommandRejected) -> do + reread <- lookupMemory mid + pure case reread of + Right (Just row) | matches row -> Right mid + _ -> Left err + other -> pure other + +memoryRecordFields :: [FieldCheck RecordMemoryData] +memoryRecordFields = + [ ("agentId", \d row -> row.agentId == d.agentId), + ("sessionId", \d row -> row.sessionId == (idText <$> d.sessionId)), + ("namespace", \d row -> row.namespace == scopeNamespaceText d.scope), + ("scopeKind", \d row -> row.scopeKind == scopeKindText d.scope), + ("scopeRef", \d row -> row.scopeRef == scopeRefText d.scope), + ("memoryType", \d row -> row.memoryType == memoryTypeToText d.memoryType), + ("content", \d row -> row.content == d.content), + ("priority", \d row -> row.priority == d.priority), + ("confidence", \d row -> row.confidence == confidenceToText d.confidence), + ("tags", \d row -> row.tags == d.tags), + ("supersedes", \d row -> row.supersedes == (idText <$> d.supersedes)) + ] + +memorySupersedeFields :: [FieldCheck SupersedeMemoryData] +memorySupersedeFields = + [ ("status", \_ row -> row.status == "superseded"), + ("supersededBy", \d row -> row.supersededBy == Just (idText d.supersededBy)) + ] + +memoryArchiveFields :: [FieldCheck ArchiveMemoryData] +memoryArchiveFields = + [("status", \_ row -> row.status == "archived")] + +-- | The projection records the merge target in @superseded_by@. Keyed on the winner id +-- alone, not a command payload, because 'merge' generates its own timestamp. +memoryMergeFields :: [FieldCheck MemoryId] +memoryMergeFields = + [ ("status", \_ row -> row.status == "merged"), + ("mergedInto", \winner row -> row.supersededBy == Just (idText winner)) + ] + +lookupMemory :: + (IOE :> es, Store :> es) => + MemoryId -> + Eff es (Either ReadModelError (Maybe MemoryRow)) +lookupMemory mid = + runQueryWith Nothing Eventual memoryByIdReadModel (MemoryByIdQuery (idText mid)) + +getMemoryRowById :: + (IOE :> es, Store :> es) => + MemoryId -> + Eff es (Either ReadModelError (Maybe MemoryRow)) +getMemoryRowById = + lookupMemory + +getActiveRowsInNamespace :: + (IOE :> es, Store :> es) => + Namespace -> + Eff es (Either ReadModelError [MemoryRow]) +getActiveRowsInNamespace (Namespace ns) = + runQueryWith Nothing Eventual memoriesByNamespaceRowsReadModel (MemoriesByNamespaceQuery ns) + +getActiveRowsByScope :: + (IOE :> es, Store :> es) => + MemoryScope -> + Eff es (Either ReadModelError [MemoryRow]) +getActiveRowsByScope scope = + runQueryWith + Nothing + Eventual + memoriesByScopeRowsReadModel + (MemoriesByScopeQuery (scopeNamespaceText scope) (scopeKindText scope) (scopeRefText scope)) + +getRowsBySession :: + (IOE :> es, Store :> es) => + SessionId -> + Eff es (Either ReadModelError [MemoryRow]) +getRowsBySession sid = + runQueryWith Nothing Eventual memoriesBySessionRowsReadModel (MemoriesBySessionQuery (idText sid)) + +getActiveRowsByType :: + (IOE :> es, Store :> es) => + Namespace -> + MemoryType -> + Eff es (Either ReadModelError [MemoryRow]) +getActiveRowsByType (Namespace ns) memoryType = + runQueryWith Nothing Eventual memoriesByTypeRowsReadModel (MemoriesByTypeQuery ns (memoryTypeToText memoryType)) + +getSupersessionChain :: + (IOE :> es, Store :> es) => + MemoryId -> + Eff es (Either ReadModelError [MemoryRow]) +getSupersessionChain mid = + runQueryWith Nothing Eventual memorySupersessionChainReadModel (MemorySupersessionChainQuery (idText mid)) + +runMemoryCommand :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + MemoryId -> + MemoryCommand -> + Eff es (Either MemoryWriteError MemoryId) +runMemoryCommand mid cmd = do + result <- + runCommandWithProjections + defaultRunCommandOptions + memoryEventStream + (memoryStream mid) + cmd + [memoryInlineProjection, l2SceneTimerScheduleProjection] + pure $ + case result of + Left err -> Left (MemoryCommandRejected err) + Right _ -> Right mid diff --git a/src/Kioku/Memory/Domain.hs b/src/Kioku/Memory/Domain.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Memory/Domain.hs @@ -0,0 +1,272 @@ +{-# LANGUAGE TemplateHaskell #-} +{-# OPTIONS_GHC -Wno-unused-top-binds #-} + +module Kioku.Memory.Domain + ( MemoryVertex (..), + MemoryRegs, + RecordMemoryData (..), + SupersedeMemoryData (..), + ArchiveMemoryData (..), + UpdateMemoryTagsData (..), + UpdateMemoryConfidenceData (..), + MergeMemoryData (..), + MemoryCommand (..), + commandMemoryId, + MemoryRecordedData (..), + MemorySupersededData (..), + MemoryArchivedData (..), + MemoryTagsUpdatedData (..), + MemoryConfidenceUpdatedData (..), + MemoryMergedData (..), + MemoryEvent (..), + eventMemoryId, + memoryTransducer, + ) +where + +import Data.Set (Set) +import Keiki.Builder qualified as B +import Keiki.Core (HsPred, SymTransducer) +import Keiki.Generics (emptyRegFile) +import Keiki.Generics.TH (deriveAggregate) +import Kioku.Api.Scope (MemoryScope) +import Kioku.Api.Types (Confidence, MemoryType) +import Kioku.Id (MemoryId, SessionId) +import Kioku.Prelude + +data MemoryVertex = NotCreated | Active | Superseded | Merged | Archived + deriving stock (Eq, Ord, Show, Enum, Bounded) + +type MemoryRegs = '[] + +data RecordMemoryData = RecordMemoryData + { memoryId :: !MemoryId, + agentId :: !Text, + sessionId :: !(Maybe SessionId), + scope :: !MemoryScope, + memoryType :: !MemoryType, + content :: !Text, + priority :: !Int, + confidence :: !Confidence, + tags :: !(Set Text), + supersedes :: !(Maybe MemoryId), + recordedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data SupersedeMemoryData = SupersedeMemoryData + { memoryId :: !MemoryId, + supersededBy :: !MemoryId, + supersededAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data ArchiveMemoryData = ArchiveMemoryData + { memoryId :: !MemoryId, + archivedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data UpdateMemoryTagsData = UpdateMemoryTagsData + { memoryId :: !MemoryId, + tags :: !(Set Text), + updatedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data UpdateMemoryConfidenceData = UpdateMemoryConfidenceData + { memoryId :: !MemoryId, + confidence :: !Confidence, + updatedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data MergeMemoryData = MergeMemoryData + { memoryId :: !MemoryId, + mergedInto :: !MemoryId, + mergedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data MemoryCommand + = RecordMemory !RecordMemoryData + | SupersedeMemory !SupersedeMemoryData + | ArchiveMemory !ArchiveMemoryData + | UpdateMemoryTags !UpdateMemoryTagsData + | UpdateMemoryConfidence !UpdateMemoryConfidenceData + | MergeMemory !MergeMemoryData + deriving stock (Generic, Eq, Show) + +commandMemoryId :: MemoryCommand -> MemoryId +commandMemoryId = \case + RecordMemory d -> d.memoryId + SupersedeMemory d -> d.memoryId + ArchiveMemory d -> d.memoryId + UpdateMemoryTags d -> d.memoryId + UpdateMemoryConfidence d -> d.memoryId + MergeMemory d -> d.memoryId + +data MemoryRecordedData = MemoryRecordedData + { memoryId :: !MemoryId, + agentId :: !Text, + sessionId :: !(Maybe SessionId), + scope :: !MemoryScope, + memoryType :: !MemoryType, + content :: !Text, + priority :: !Int, + confidence :: !Confidence, + tags :: !(Set Text), + supersedes :: !(Maybe MemoryId), + recordedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data MemorySupersededData = MemorySupersededData + { memoryId :: !MemoryId, + supersededBy :: !MemoryId, + supersededAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data MemoryArchivedData = MemoryArchivedData + { memoryId :: !MemoryId, + archivedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data MemoryTagsUpdatedData = MemoryTagsUpdatedData + { memoryId :: !MemoryId, + tags :: !(Set Text), + updatedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data MemoryConfidenceUpdatedData = MemoryConfidenceUpdatedData + { memoryId :: !MemoryId, + confidence :: !Confidence, + updatedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data MemoryMergedData = MemoryMergedData + { memoryId :: !MemoryId, + mergedInto :: !MemoryId, + mergedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data MemoryEvent + = MemoryRecorded !MemoryRecordedData + | MemorySuperseded !MemorySupersededData + | MemoryArchived !MemoryArchivedData + | MemoryTagsUpdated !MemoryTagsUpdatedData + | MemoryConfidenceUpdated !MemoryConfidenceUpdatedData + | MemoryMerged !MemoryMergedData + deriving stock (Generic, Eq, Show) + +instance FromJSON MemoryEvent where + parseJSON = genericParseJSON eventAesonOptions + +instance ToJSON MemoryEvent where + toJSON = genericToJSON eventAesonOptions + +eventMemoryId :: MemoryEvent -> MemoryId +eventMemoryId = \case + MemoryRecorded d -> d.memoryId + MemorySuperseded d -> d.memoryId + MemoryArchived d -> d.memoryId + MemoryTagsUpdated d -> d.memoryId + MemoryConfidenceUpdated d -> d.memoryId + MemoryMerged d -> d.memoryId + +$(deriveAggregate ''MemoryCommand ''MemoryRegs ''MemoryEvent) + +memoryTransducer :: + SymTransducer + (HsPred MemoryRegs MemoryCommand) + MemoryRegs + MemoryVertex + MemoryCommand + MemoryEvent +memoryTransducer = + B.buildTransducer NotCreated emptyRegFile isTerminal do + B.from NotCreated do + B.onCmd inCtorRecordMemory $ \d -> B.do + B.emit + wireMemoryRecorded + MemoryRecordedTermFields + { memoryId = d.memoryId, + agentId = d.agentId, + sessionId = d.sessionId, + scope = d.scope, + memoryType = d.memoryType, + content = d.content, + priority = d.priority, + confidence = d.confidence, + tags = d.tags, + supersedes = d.supersedes, + recordedAt = d.recordedAt + } + B.goto Active + + B.from Active do + B.onCmd inCtorSupersedeMemory $ \d -> B.do + B.emit + wireMemorySuperseded + MemorySupersededTermFields + { memoryId = d.memoryId, + supersededBy = d.supersededBy, + supersededAt = d.supersededAt + } + B.goto Superseded + + B.onCmd inCtorArchiveMemory $ \d -> B.do + B.emit + wireMemoryArchived + MemoryArchivedTermFields + { memoryId = d.memoryId, + archivedAt = d.archivedAt + } + B.goto Archived + + B.onCmd inCtorUpdateMemoryTags $ \d -> B.do + B.emit + wireMemoryTagsUpdated + MemoryTagsUpdatedTermFields + { memoryId = d.memoryId, + tags = d.tags, + updatedAt = d.updatedAt + } + B.goto Active + + B.onCmd inCtorUpdateMemoryConfidence $ \d -> B.do + B.emit + wireMemoryConfidenceUpdated + MemoryConfidenceUpdatedTermFields + { memoryId = d.memoryId, + confidence = d.confidence, + updatedAt = d.updatedAt + } + B.goto Active + + B.onCmd inCtorMergeMemory $ \d -> B.do + B.emit + wireMemoryMerged + MemoryMergedTermFields + { memoryId = d.memoryId, + mergedInto = d.mergedInto, + mergedAt = d.mergedAt + } + B.goto Merged + where + isTerminal = \case + Superseded -> True + Merged -> True + Archived -> True + _ -> False diff --git a/src/Kioku/Memory/Embedding.hs b/src/Kioku/Memory/Embedding.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Memory/Embedding.hs @@ -0,0 +1,95 @@ +module Kioku.Memory.Embedding + ( EmbeddingConfig (..), + EmbedError (..), + resolveEmbeddingConfig, + toEmbeddingModel, + embedWithRetry, + sha256Hex, + ) +where + +import Baikai.Auth (ApiKeySource (..)) +import Baikai.Embedding (EmbeddingModel (..), embedOne) +import Control.Concurrent (threadDelay) +import Control.Exception (SomeException, try) +import Crypto.Hash (Digest, SHA256) +import Crypto.Hash qualified as Hash +import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import Kioku.Prelude +import Numeric.Natural (Natural) +import System.Environment (lookupEnv) +import Text.Read (readMaybe) + +data EmbeddingConfig = EmbeddingConfig + { baseUrl :: !Text, + model :: !Text, + dimensions :: !Int, + apiKey :: !Text + } + deriving stock (Generic, Eq, Show) + +data EmbedError + = EmbedTransport !Text + | EmbedEmpty + deriving stock (Generic, Eq, Show) + +resolveEmbeddingConfig :: IO EmbeddingConfig +resolveEmbeddingConfig = do + baseUrl <- envText "KIOKU_EMBEDDING_BASE_URL" "https://api.openai.com" + model <- envText "KIOKU_EMBEDDING_MODEL" "text-embedding-3-small" + dimensions <- envInt "KIOKU_EMBEDDING_DIMENSIONS" 1536 + apiKey <- envTextFallback ["KIOKU_EMBEDDING_API_KEY", "OPENAI_API_KEY"] "" + pure EmbeddingConfig {baseUrl, model, dimensions, apiKey} + +toEmbeddingModel :: EmbeddingConfig -> EmbeddingModel +toEmbeddingModel cfg = + EmbeddingModel + { modelId = cfg.model, + baseUrl = cfg.baseUrl, + dimensions = Just (fromIntegral @Int @Natural cfg.dimensions), + apiKey = ApiKeyLiteral cfg.apiKey + } + +embedWithRetry :: EmbeddingModel -> Int -> Text -> IO (Either EmbedError (Vector Double)) +embedWithRetry model maxAttempts input = go 1 + where + attempts = max 1 maxAttempts + + go attempt = do + result <- try (embedOne model input) + case result of + Right vec + | Vector.null vec -> pure (Left EmbedEmpty) + | otherwise -> pure (Right vec) + Left (err :: SomeException) + | attempt >= attempts -> pure (Left (EmbedTransport (Text.pack (show err)))) + | otherwise -> do + threadDelay (attemptDelayMicros attempt) + go (attempt + 1) + +sha256Hex :: Text -> Text +sha256Hex content = + Text.pack (show (Hash.hash (TE.encodeUtf8 content) :: Digest SHA256)) + +attemptDelayMicros :: Int -> Int +attemptDelayMicros attempt = 200000 * (2 ^ max 0 (attempt - 1)) + +envText :: String -> Text -> IO Text +envText name fallback = + maybe fallback Text.pack <$> lookupEnv name + +envInt :: String -> Int -> IO Int +envInt name fallback = do + found <- lookupEnv name + pure $ fromMaybe fallback (found >>= readMaybe) + +envTextFallback :: [String] -> Text -> IO Text +envTextFallback [] fallback = pure fallback +envTextFallback (name : rest) fallback = do + found <- lookupEnv name + case found of + Just value -> pure (Text.pack value) + Nothing -> envTextFallback rest fallback diff --git a/src/Kioku/Memory/Embedding/Worker.hs b/src/Kioku/Memory/Embedding/Worker.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Memory/Embedding/Worker.hs @@ -0,0 +1,368 @@ +module Kioku.Memory.Embedding.Worker + ( EmbeddingWorkerEnv (..), + EmbedOutcome (..), + backfillMissingEmbeddings, + embeddingHandler, + embeddingWorkerProcessor, + mkEmbeddingWorkerEnv, + runEmbeddingWorkerHost, + shouldSkipEmbedding, + ) +where + +import Baikai.Embedding (EmbeddingModel (..)) +import Control.Monad (foldM) +import Data.Functor.Contravariant ((>$<)) +import Data.Int (Int32) +import Data.Set qualified as Set +import Data.Text qualified as Text +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import Effectful (Eff, IOE, (:>)) +import Effectful.Error.Static (Error) +import Effectful.Error.Static qualified as EffError +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Keiro.Codec (decodeRecorded) +import Kioku.Id (MemoryId, idText) +import Kioku.Memory.Domain (MemoryEvent (..), MemoryRecordedData (..)) +import Kioku.Memory.Embedding (EmbedError, embedWithRetry, sha256Hex) +import Kioku.Memory.EventStream (memoryCodec) +import Kioku.Prelude +import Kioku.Recall.Capability (VectorCapability (..)) +import Kioku.Worker.Failure (embeddingRetryDelay, isTransientStoreError) +import Kiroku.Store.Connection (KirokuStore) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Transaction (runTransaction) +import Kiroku.Store.Types (CategoryName (..), EventType (..), RecordedEvent) +import Shibuya.Adapter.Kiroku + ( EventTypeFilter (..), + KirokuAdapterConfig (..), + SubscriptionName (..), + SubscriptionTarget (..), + defaultKirokuAdapterConfig, + guardKirokuHandler, + kirokuAdapter, + ) +import Shibuya.App (ProcessorId (..), QueueProcessor (..), defaultAppConfig, runApp, waitApp) +import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..), HaltReason (..)) +import Shibuya.Core.Ingested (Ingested (..), Message (..)) +import Shibuya.Core.Types (Envelope (..)) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..)) +import Shibuya.Telemetry.Effect (Tracing) +import System.IO qualified as IO + +data EmbeddingCandidate = EmbeddingCandidate + { memoryId :: !Text, + content :: !Text, + contentHash :: !(Maybe Text), + hasEmbedding :: !Bool + } + deriving stock (Generic, Eq, Show) + +data EmbeddingUpdate = EmbeddingUpdate + { memoryId :: !Text, + embedding :: !(Vector Double), + embeddingModel :: !Text, + dimensions :: !Int, + contentHash :: !Text + } + deriving stock (Generic, Eq, Show) + +data EmbeddingState = EmbeddingState + { contentHash :: !(Maybe Text), + hasEmbedding :: !Bool + } + deriving stock (Generic, Eq, Show) + +-- | Everything the embedding path needs from the outside world. +-- +-- The provider call is a field rather than a direct 'embedWithRetry' call so +-- tests can drive every branch of the ack taxonomy — a failing provider, a +-- succeeding one, one that returns the wrong number of dimensions — without an +-- embedding API key or a network. +data EmbeddingWorkerEnv = EmbeddingWorkerEnv + { model :: !EmbeddingModel, + dimensions :: !Int, + embed :: !(Text -> IO (Either EmbedError (Vector Double))) + } + deriving stock (Generic) + +-- | The production environment: the real provider, retried three times +-- in-process (~0.6s of jitter-free backoff) before the failure is reported to +-- the caller, which then decides whether the /event/ should be redelivered. +mkEmbeddingWorkerEnv :: EmbeddingModel -> Int -> EmbeddingWorkerEnv +mkEmbeddingWorkerEnv model dims = + EmbeddingWorkerEnv {model, dimensions = dims, embed = embedWithRetry model 3} + +-- | What one embedding attempt did. +-- +-- 'EmbedSkipped' covers both "already embedded with this exact content" and +-- "the memory is not there to embed"; neither is a failure. The distinction +-- that matters to the handler is 'EmbedFailed', which used to be indistinguishable +-- from success. +data EmbedOutcome + = EmbedSkipped + | EmbedStored + | EmbedFailed !EmbedError + deriving stock (Generic, Eq, Show) + +runEmbeddingWorkerHost :: + (IOE :> es, Store :> es, Error StoreError :> es, Tracing :> es) => + KirokuStore -> + VectorCapability -> + EmbeddingModel -> + Int -> + Eff es () +runEmbeddingWorkerHost store capability model dims = do + processor <- embeddingWorkerProcessor capability model dims store + started <- runApp defaultAppConfig [processor] + case started of + Left appErr -> + liftIO (ioError (userError ("kioku embedding worker failed to start: " <> show appErr))) + Right appHandle -> do + liftIO (putStrLn "kioku embedding worker started. Press Ctrl+C to stop.") + waitApp appHandle + +embeddingWorkerProcessor :: + (IOE :> es, Store :> es, Error StoreError :> es) => + VectorCapability -> + EmbeddingModel -> + Int -> + KirokuStore -> + Eff es (ProcessorId, QueueProcessor es) +embeddingWorkerProcessor capability model dims store = do + adapter <- kirokuAdapter store embeddingAdapterConfig + pure + ( ProcessorId embeddingWorkerName, + QueueProcessor + { adapter, + -- The kiroku bridge is ack-coupled: a synchronous exception escaping + -- the handler leaves the ack unfinalized and blocks the subscription + -- worker forever. The guard turns that into a one-second retry. + handler = guardKirokuHandler (embeddingMessageHandler capability (mkEmbeddingWorkerEnv model dims)), + ordering = StrictInOrder, + concurrency = Serial + } + ) + +-- | Decide what happens to one delivered @MemoryRecorded@ event. +-- +-- Every branch is a deliberate choice about durability: +-- +-- * a provider failure is /transient/ — retry with backoff, and let kiroku's +-- retry policy dead-letter it if the outage outlasts the window; +-- * an undecodable payload can never succeed — dead-letter it visibly rather +-- than acking it into the void; +-- * a transient store error must not kill the pipeline — retry; +-- * a permanent store error (a dimension mismatch, a broken schema) would fail +-- identically for every subsequent event — halting is the honest response, +-- because dead-lettering would quietly drain the whole stream. +embeddingHandler :: + (IOE :> es, Store :> es, Error StoreError :> es) => + VectorCapability -> + EmbeddingWorkerEnv -> + Ingested es RecordedEvent -> + Eff es AckDecision +embeddingHandler capability env ingested = + handleEmbeddingEnvelope capability env ingested.envelope + +embeddingMessageHandler :: + (IOE :> es, Store :> es, Error StoreError :> es) => + VectorCapability -> + EmbeddingWorkerEnv -> + Message es RecordedEvent -> + Eff es AckDecision +embeddingMessageHandler capability env message = + handleEmbeddingEnvelope capability env message.envelope + +handleEmbeddingEnvelope :: + (IOE :> es, Store :> es, Error StoreError :> es) => + VectorCapability -> + EmbeddingWorkerEnv -> + Envelope RecordedEvent -> + Eff es AckDecision +handleEmbeddingEnvelope capability env envelope = + EffError.catchError @StoreError run \_callStack storeErr -> + if isTransientStoreError storeErr + then do + logWorker ("transient store error, retrying: " <> Text.pack (show storeErr)) + pure (AckRetry retryDelay) + else do + logWorker ("fatal store error, halting: " <> Text.pack (show storeErr)) + pure (AckHalt (HaltFatal ("kioku embedding worker store error: " <> Text.pack (show storeErr)))) + where + retryDelay = embeddingRetryDelay envelope.attempt + + run = + case decodeRecorded memoryCodec envelope.payload of + Left codecErr -> do + logWorker ("undecodable event, dead-lettering: " <> Text.pack (show codecErr)) + pure (AckDeadLetter (InvalidPayload (Text.pack (show codecErr)))) + Right (MemoryRecorded d) -> do + outcome <- embedMemoryContent capability env (idText (d.memoryId :: MemoryId)) d.content + case outcome of + EmbedFailed err -> do + logWorker ("embedding failed, retrying: " <> Text.pack (show err)) + pure (AckRetry retryDelay) + EmbedStored -> pure AckOk + EmbedSkipped -> pure AckOk + -- The subscription is filtered to MemoryRecorded, so this is unreachable + -- today; acking is the harmless answer if the filter ever widens. + Right _ -> pure AckOk + +logWorker :: (IOE :> es) => Text -> Eff es () +logWorker msg = + liftIO (IO.hPutStrLn IO.stderr (Text.unpack (embeddingWorkerName <> ": " <> msg))) + +backfillMissingEmbeddings :: + (IOE :> es, Store :> es) => + VectorCapability -> + EmbeddingModel -> + Int -> + Eff es Int +backfillMissingEmbeddings VectorAvailable model dims = do + candidates <- runTransaction (Tx.statement () selectEmbeddingCandidatesStmt) + foldM embedCandidate 0 candidates + where + env = mkEmbeddingWorkerEnv model dims + + embedCandidate count candidate + | shouldSkipEmbedding candidate.hasEmbedding candidate.contentHash contentHash = + pure count + | otherwise = do + outcome <- embedAndStore env candidate.memoryId candidate.content contentHash + case outcome of + EmbedStored -> pure (count + 1) + EmbedSkipped -> pure count + -- One unembeddable memory must not abort the pass: a backfill exists + -- precisely to recover from failures, and the next run retries this row. + EmbedFailed err -> do + logWorker ("backfill skipped " <> candidate.memoryId <> ": " <> Text.pack (show err)) + pure count + where + contentHash = sha256Hex candidate.content +backfillMissingEmbeddings _ _ _ = pure 0 + +embedMemoryContent :: + (IOE :> es, Store :> es) => + VectorCapability -> + EmbeddingWorkerEnv -> + Text -> + Text -> + Eff es EmbedOutcome +embedMemoryContent VectorAvailable env memoryId content = do + existing <- runTransaction (Tx.statement memoryId selectEmbeddingStateStmt) + case existing of + Nothing -> pure EmbedSkipped + Just state + | shouldSkipEmbedding state.hasEmbedding state.contentHash contentHash -> + pure EmbedSkipped + | otherwise -> + embedAndStore env memoryId content contentHash + where + contentHash = sha256Hex content +embedMemoryContent _ _ _ _ = pure EmbedSkipped + +shouldSkipEmbedding :: Bool -> Maybe Text -> Text -> Bool +shouldSkipEmbedding hasEmbedding storedContentHash contentHash = + hasEmbedding && storedContentHash == Just contentHash + +embedAndStore :: + (IOE :> es, Store :> es) => + EmbeddingWorkerEnv -> + Text -> + Text -> + Text -> + Eff es EmbedOutcome +embedAndStore env memoryId content contentHash = do + result <- liftIO (env.embed content) + case result of + Left err -> pure (EmbedFailed err) + Right embedding -> do + runTransaction $ + Tx.statement + EmbeddingUpdate + { memoryId, + embedding, + embeddingModel = env.model.modelId, + dimensions = env.dimensions, + contentHash + } + upsertEmbeddingStmt + pure EmbedStored + +selectEmbeddingCandidatesStmt :: Statement () [EmbeddingCandidate] +selectEmbeddingCandidatesStmt = + preparable + """ + SELECT memory_id, content, content_hash, embedding IS NOT NULL AS has_embedding + FROM kiroku.kioku_memories + WHERE status = 'active' + ORDER BY created_at ASC + """ + E.noParams + (D.rowList embeddingCandidateDecoder) + +selectEmbeddingStateStmt :: Statement Text (Maybe EmbeddingState) +selectEmbeddingStateStmt = + preparable + """ + SELECT content_hash, embedding IS NOT NULL AS has_embedding + FROM kiroku.kioku_memories + WHERE memory_id = $1 AND status = 'active' + """ + (E.param (E.nonNullable E.text)) + (D.rowMaybe embeddingStateDecoder) + +embeddingCandidateDecoder :: D.Row EmbeddingCandidate +embeddingCandidateDecoder = + EmbeddingCandidate + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nonNullable D.bool) + +embeddingStateDecoder :: D.Row EmbeddingState +embeddingStateDecoder = + EmbeddingState + <$> D.column (D.nullable D.text) + <*> D.column (D.nonNullable D.bool) + +upsertEmbeddingStmt :: Statement EmbeddingUpdate () +upsertEmbeddingStmt = + preparable + """ + UPDATE kiroku.kioku_memories + SET embedding = $2::vector, + embedding_model = $3, + dimensions = $4, + content_hash = $5 + WHERE memory_id = $1 + """ + embeddingUpdateEncoder + D.noResult + +embeddingUpdateEncoder :: E.Params EmbeddingUpdate +embeddingUpdateEncoder = + ((\update -> update.memoryId) >$< E.param (E.nonNullable E.text)) + <> ((\update -> vectorLiteral update.embedding) >$< E.param (E.nonNullable E.text)) + <> ((\update -> update.embeddingModel) >$< E.param (E.nonNullable E.text)) + <> ((\update -> fromIntegral @Int @Int32 update.dimensions) >$< E.param (E.nonNullable E.int4)) + <> ((\update -> update.contentHash) >$< E.param (E.nonNullable E.text)) + +vectorLiteral :: Vector Double -> Text +vectorLiteral values = + "[" <> Text.intercalate "," (Text.pack . show <$> Vector.toList values) <> "]" + +embeddingAdapterConfig :: KirokuAdapterConfig +embeddingAdapterConfig = + (defaultKirokuAdapterConfig (SubscriptionName embeddingWorkerName) (Category (CategoryName "kioku_memory"))) + { eventTypeFilter = OnlyEventTypes (Set.fromList [EventType "MemoryRecorded"]) + } + +embeddingWorkerName :: Text +embeddingWorkerName = "kioku-memory-embedding" diff --git a/src/Kioku/Memory/EventStream.hs b/src/Kioku/Memory/EventStream.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Memory/EventStream.hs @@ -0,0 +1,162 @@ +module Kioku.Memory.EventStream + ( MemoryEventStream, + memoryEventStream, + memoryCodec, + memoryStream, + parseMemoryEvent, + ) +where + +import Data.Aeson (Value) +import Data.Aeson.Types (Parser, parseEither, withObject, (.:), (.:?)) +import Data.Text qualified as Text +import Keiki.Core (HsPred) +import Keiki.Generics (emptyRegFile) +import Keiro.Codec (Codec (..), EventType (..)) +import Keiro.EventStream (EventStream (..), SnapshotPolicy (..)) +import Keiro.EventStream.Validate (ValidatedEventStream, mkEventStreamOrThrow) +import Keiro.Stream (Stream) +import Keiro.Stream qualified as Stream +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), ScopeKind (..)) +import Kioku.Id (MemoryId, SessionId, idText, parseIdLenient) +import Kioku.Memory.Domain +import Kioku.Prelude + +type MemoryEventStream = + EventStream (HsPred MemoryRegs MemoryCommand) MemoryRegs MemoryVertex MemoryCommand MemoryEvent + +memoryStream :: MemoryId -> Stream MemoryEventStream +memoryStream mid = Stream.entityStream (Stream.categoryUnsafe "kioku_memory") (idText mid) + +memoryEventStream :: ValidatedEventStream (HsPred MemoryRegs MemoryCommand) MemoryRegs MemoryVertex MemoryCommand MemoryEvent +memoryEventStream = + mkEventStreamOrThrow "kioku-memory" memoryEventStreamDefinition + +memoryEventStreamDefinition :: MemoryEventStream +memoryEventStreamDefinition = + EventStream + { transducer = memoryTransducer, + initialState = NotCreated, + initialRegisters = emptyRegFile, + eventCodec = memoryCodec, + resolveStreamName = Stream.streamName, + snapshotPolicy = Never, + stateCodec = Nothing + } + +memoryCodec :: Codec MemoryEvent +memoryCodec = + Codec + { eventTypes = + EventType + <$> "MemoryRecorded" + :| [ "MemorySuperseded", + "MemoryArchived", + "MemoryTagsUpdated", + "MemoryConfidenceUpdated", + "MemoryMerged" + ], + eventType = + EventType . \case + MemoryRecorded {} -> "MemoryRecorded" + MemorySuperseded {} -> "MemorySuperseded" + MemoryArchived {} -> "MemoryArchived" + MemoryTagsUpdated {} -> "MemoryTagsUpdated" + MemoryConfidenceUpdated {} -> "MemoryConfidenceUpdated" + MemoryMerged {} -> "MemoryMerged", + schemaVersion = 1, + encode = toJSON, + decode = const parseMemoryEvent, + upcasters = [] + } + +parseMemoryEvent :: Value -> Either Text MemoryEvent +parseMemoryEvent value = + case parseEither parseJSON value of + Right event -> Right event + Left nativeErr -> + case parseEither parseLegacyMemoryEvent value of + Right event -> Right event + Left legacyErr -> Left (Text.pack nativeErr <> "; legacy decode failed: " <> Text.pack legacyErr) + +parseLegacyMemoryEvent :: Value -> Parser MemoryEvent +parseLegacyMemoryEvent = + withObject "Rei AgentMemoryEvent" $ \o -> do + tag <- o .: "type" + payload <- o .: "data" + case tag of + "agent_memory_recorded" -> MemoryRecorded <$> parseLegacyMemoryRecorded payload + "agent_memory_superseded" -> MemorySuperseded <$> parseLegacyMemorySuperseded payload + "agent_memory_archived" -> MemoryArchived <$> parseLegacyMemoryArchived payload + "agent_memory_tags_updated" -> MemoryTagsUpdated <$> parseLegacyMemoryTagsUpdated payload + "agent_memory_confidence_updated" -> MemoryConfidenceUpdated <$> parseLegacyMemoryConfidenceUpdated payload + other -> fail ("Unknown Rei AgentMemoryEvent tag: " <> Text.unpack other) + +parseLegacyMemoryRecorded :: Value -> Parser MemoryRecordedData +parseLegacyMemoryRecorded = + withObject "Rei AgentMemoryRecordedData" $ \o -> do + memoryId <- parseLegacyMemoryId =<< o .: "memoryId" + sessionId <- traverse parseLegacySessionId =<< o .:? "sessionId" + scope <- parseLegacyAnchor =<< o .: "anchor" + supersedes <- traverse parseLegacyMemoryId =<< o .:? "supersedes" + MemoryRecordedData memoryId + <$> o .: "agentId" + <*> pure sessionId + <*> pure scope + <*> o .: "memoryType" + <*> o .: "content" + <*> pure 100 + <*> o .: "confidence" + <*> o .: "tags" + <*> pure supersedes + <*> o .: "recordedAt" + +parseLegacyMemorySuperseded :: Value -> Parser MemorySupersededData +parseLegacyMemorySuperseded = + withObject "Rei AgentMemorySupersededData" $ \o -> + MemorySupersededData + <$> (parseLegacyMemoryId =<< o .: "memoryId") + <*> (parseLegacyMemoryId =<< o .: "supersededBy") + <*> o .: "supersededAt" + +parseLegacyMemoryArchived :: Value -> Parser MemoryArchivedData +parseLegacyMemoryArchived = + withObject "Rei AgentMemoryArchivedData" $ \o -> + MemoryArchivedData + <$> (parseLegacyMemoryId =<< o .: "memoryId") + <*> o .: "archivedAt" + +parseLegacyMemoryTagsUpdated :: Value -> Parser MemoryTagsUpdatedData +parseLegacyMemoryTagsUpdated = + withObject "Rei AgentMemoryTagsUpdatedData" $ \o -> + MemoryTagsUpdatedData + <$> (parseLegacyMemoryId =<< o .: "memoryId") + <*> o .: "tags" + <*> o .: "updatedAt" + +parseLegacyMemoryConfidenceUpdated :: Value -> Parser MemoryConfidenceUpdatedData +parseLegacyMemoryConfidenceUpdated = + withObject "Rei AgentMemoryConfidenceUpdatedData" $ \o -> + MemoryConfidenceUpdatedData + <$> (parseLegacyMemoryId =<< o .: "memoryId") + <*> o .: "confidence" + <*> o .: "updatedAt" + +parseLegacyAnchor :: Value -> Parser MemoryScope +parseLegacyAnchor = + withObject "Rei MemoryAnchor" $ \o -> do + anchorType <- o .: "type" + case anchorType of + "intention" -> ScopeEntity reiNamespace (ScopeKind "intention") <$> o .: "id" + "habit" -> ScopeEntity reiNamespace (ScopeKind "habit") <$> o .: "id" + "workspace" -> pure (ScopeGlobal reiNamespace) + other -> fail ("Unknown Rei MemoryAnchor type: " <> Text.unpack other) + +parseLegacyMemoryId :: Text -> Parser MemoryId +parseLegacyMemoryId = either (fail . Text.unpack) pure . parseIdLenient + +parseLegacySessionId :: Text -> Parser SessionId +parseLegacySessionId = either (fail . Text.unpack) pure . parseIdLenient + +reiNamespace :: Namespace +reiNamespace = Namespace "rei" diff --git a/src/Kioku/Memory/ReadModel.hs b/src/Kioku/Memory/ReadModel.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Memory/ReadModel.hs @@ -0,0 +1,575 @@ +module Kioku.Memory.ReadModel + ( memoryInlineProjection, + MemoryRow (..), + MemoryByIdQuery (..), + MemoriesByNamespaceQuery (..), + MemoriesByScopeQuery (..), + MemoriesBySessionQuery (..), + MemoriesByTypeQuery (..), + MemorySupersessionChainQuery (..), + memoryByIdReadModel, + memoriesByNamespaceReadModel, + memoriesByNamespaceRowsReadModel, + memoriesByScopeReadModel, + memoriesByScopeRowsReadModel, + memoriesBySessionReadModel, + memoriesBySessionRowsReadModel, + memoriesByTypeReadModel, + memoriesByTypeRowsReadModel, + memorySupersessionChainReadModel, + ) +where + +import Contravariant.Extras (contrazip2, contrazip3) +import Data.Aeson qualified as Aeson +import Data.ByteString.Lazy qualified as BL +import Data.Functor.Contravariant ((>$<)) +import Data.Generics.Labels () +import Data.Int (Int32) +import Data.Set (Set) +import Data.Set qualified as Set +import Data.Text.Encoding qualified as TE +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Keiro.Projection (InlineProjection (..)) +import Keiro.ReadModel (ConsistencyMode (..), ReadModel (..), StrongScope (..)) +import Kioku.Api.Scope (scopeFromColumns, scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Api.Types (MemoryRecord (..), confidenceToText, memoryTypeToText) +import Kioku.Id (idText) +import Kioku.Memory.Domain +import Kioku.Prelude +import Kiroku.Store.Types (RecordedEvent) + +data MemoryRow = MemoryRow + { memoryId :: !Text, + agentId :: !Text, + sessionId :: !(Maybe Text), + namespace :: !Text, + scopeKind :: !(Maybe Text), + scopeRef :: !(Maybe Text), + memoryType :: !Text, + content :: !Text, + priority :: !Int, + confidence :: !Text, + tags :: !(Set Text), + status :: !Text, + supersededBy :: !(Maybe Text), + supersedes :: !(Maybe Text), + createdAt :: !UTCTime, + updatedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +newtype MemoryByIdQuery = MemoryByIdQuery Text + +newtype MemoriesByNamespaceQuery = MemoriesByNamespaceQuery Text + +data MemoriesByScopeQuery = MemoriesByScopeQuery Text (Maybe Text) (Maybe Text) + +newtype MemoriesBySessionQuery = MemoriesBySessionQuery Text + +data MemoriesByTypeQuery = MemoriesByTypeQuery Text Text + +newtype MemorySupersessionChainQuery = MemorySupersessionChainQuery Text + +memoryInlineProjection :: InlineProjection MemoryEvent +memoryInlineProjection = + InlineProjection + { name = "kioku-memory-inline", + apply = applyMemoryEvent + } + +applyMemoryEvent :: MemoryEvent -> RecordedEvent -> Tx.Transaction () +applyMemoryEvent event _recorded = + case event of + MemoryRecorded d -> Tx.statement (recordedRow d) upsertMemoryStmt + MemorySuperseded d -> + Tx.statement + (idText d.memoryId, idText d.supersededBy, d.supersededAt) + updateMemorySupersededStmt + MemoryArchived d -> + Tx.statement (idText d.memoryId, d.archivedAt) updateMemoryArchivedStmt + MemoryTagsUpdated d -> + Tx.statement (idText d.memoryId, d.tags, d.updatedAt) updateMemoryTagsStmt + MemoryConfidenceUpdated d -> + Tx.statement (idText d.memoryId, confidenceToText d.confidence, d.updatedAt) updateMemoryConfidenceStmt + MemoryMerged d -> + Tx.statement (idText d.memoryId, idText d.mergedInto, d.mergedAt) updateMemoryMergedStmt + +recordedRow :: MemoryRecordedData -> MemoryRow +recordedRow d = + MemoryRow + { memoryId = idText d.memoryId, + agentId = d.agentId, + sessionId = idText <$> d.sessionId, + namespace = scopeNamespaceText d.scope, + scopeKind = scopeKindText d.scope, + scopeRef = scopeRefText d.scope, + memoryType = memoryTypeToText d.memoryType, + content = d.content, + priority = d.priority, + confidence = confidenceToText d.confidence, + tags = d.tags, + status = "active", + supersededBy = Nothing, + supersedes = idText <$> d.supersedes, + createdAt = d.recordedAt, + updatedAt = d.recordedAt + } + +memoryByIdReadModel :: ReadModel MemoryByIdQuery (Maybe MemoryRow) +memoryByIdReadModel = + ReadModel + { name = "kioku-memory-by-id", + schema = "kiroku", + tableName = "kioku_memories", + subscriptionName = "kioku-memory-inline", + version = 1, + shapeHash = "kioku-memory-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(MemoryByIdQuery mid) -> Tx.statement mid selectMemoryByIdStmt + } + +memoriesByNamespaceReadModel :: ReadModel MemoriesByNamespaceQuery [MemoryRecord] +memoriesByNamespaceReadModel = + ReadModel + { name = "kioku-memories-by-namespace", + schema = "kiroku", + tableName = "kioku_memories", + subscriptionName = "kioku-memory-inline", + version = 1, + shapeHash = "kioku-memory-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(MemoriesByNamespaceQuery ns) -> Tx.statement ns selectActiveByNamespaceStmt + } + +memoriesByNamespaceRowsReadModel :: ReadModel MemoriesByNamespaceQuery [MemoryRow] +memoriesByNamespaceRowsReadModel = + ReadModel + { name = "kioku-memory-rows-by-namespace", + schema = "kiroku", + tableName = "kioku_memories", + subscriptionName = "kioku-memory-inline", + version = 1, + shapeHash = "kioku-memory-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(MemoriesByNamespaceQuery ns) -> Tx.statement ns selectActiveByNamespaceRowsStmt + } + +memoriesByScopeReadModel :: ReadModel MemoriesByScopeQuery [MemoryRecord] +memoriesByScopeReadModel = + ReadModel + { name = "kioku-memories-by-scope", + schema = "kiroku", + tableName = "kioku_memories", + subscriptionName = "kioku-memory-inline", + version = 1, + shapeHash = "kioku-memory-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(MemoriesByScopeQuery ns sk sr) -> Tx.statement (ns, sk, sr) selectActiveByScopeStmt + } + +memoriesByScopeRowsReadModel :: ReadModel MemoriesByScopeQuery [MemoryRow] +memoriesByScopeRowsReadModel = + ReadModel + { name = "kioku-memory-rows-by-scope", + schema = "kiroku", + tableName = "kioku_memories", + subscriptionName = "kioku-memory-inline", + version = 1, + shapeHash = "kioku-memory-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(MemoriesByScopeQuery ns sk sr) -> Tx.statement (ns, sk, sr) selectActiveByScopeRowsStmt + } + +memoriesBySessionReadModel :: ReadModel MemoriesBySessionQuery [MemoryRecord] +memoriesBySessionReadModel = + ReadModel + { name = "kioku-memories-by-session", + schema = "kiroku", + tableName = "kioku_memories", + subscriptionName = "kioku-memory-inline", + version = 1, + shapeHash = "kioku-memory-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(MemoriesBySessionQuery sid) -> Tx.statement sid selectBySessionStmt + } + +memoriesBySessionRowsReadModel :: ReadModel MemoriesBySessionQuery [MemoryRow] +memoriesBySessionRowsReadModel = + ReadModel + { name = "kioku-memory-rows-by-session", + schema = "kiroku", + tableName = "kioku_memories", + subscriptionName = "kioku-memory-inline", + version = 1, + shapeHash = "kioku-memory-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(MemoriesBySessionQuery sid) -> Tx.statement sid selectBySessionRowsStmt + } + +memoriesByTypeReadModel :: ReadModel MemoriesByTypeQuery [MemoryRecord] +memoriesByTypeReadModel = + ReadModel + { name = "kioku-memories-by-type", + schema = "kiroku", + tableName = "kioku_memories", + subscriptionName = "kioku-memory-inline", + version = 1, + shapeHash = "kioku-memory-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(MemoriesByTypeQuery ns mt) -> Tx.statement (ns, mt) selectByTypeStmt + } + +memoriesByTypeRowsReadModel :: ReadModel MemoriesByTypeQuery [MemoryRow] +memoriesByTypeRowsReadModel = + ReadModel + { name = "kioku-memory-rows-by-type", + schema = "kiroku", + tableName = "kioku_memories", + subscriptionName = "kioku-memory-inline", + version = 1, + shapeHash = "kioku-memory-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(MemoriesByTypeQuery ns mt) -> Tx.statement (ns, mt) selectByTypeRowsStmt + } + +memorySupersessionChainReadModel :: ReadModel MemorySupersessionChainQuery [MemoryRow] +memorySupersessionChainReadModel = + ReadModel + { name = "kioku-memory-supersession-chain", + schema = "kiroku", + tableName = "kioku_memories", + subscriptionName = "kioku-memory-inline", + version = 1, + shapeHash = "kioku-memory-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(MemorySupersessionChainQuery mid) -> Tx.statement mid selectSupersessionChainStmt + } + +memoryRowDecoder :: D.Row MemoryRow +memoryRowDecoder = + MemoryRow + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> (fromIntegral @Int32 @Int <$> D.column (D.nonNullable D.int4)) + <*> D.column (D.nonNullable D.text) + <*> (decodeTags <$> D.column (D.nonNullable D.text)) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nonNullable D.timestamptz) + <*> D.column (D.nonNullable D.timestamptz) + +memoryRecordDecoder :: D.Row MemoryRecord +memoryRecordDecoder = + toRecord <$> memoryRowDecoder + +toRecord :: MemoryRow -> MemoryRecord +toRecord row = + MemoryRecord + { memoryId = row.memoryId, + agentId = row.agentId, + sessionId = row.sessionId, + scope = scopeFromColumns row.namespace row.scopeKind row.scopeRef, + memoryType = row.memoryType, + content = row.content, + priority = row.priority, + confidence = row.confidence, + tags = row.tags, + status = row.status, + createdAt = row.createdAt + } + +encodeTags :: Set Text -> Text +encodeTags = TE.decodeUtf8 . BL.toStrict . Aeson.encode + +decodeTags :: Text -> Set Text +decodeTags = + fromMaybe Set.empty . Aeson.decode . BL.fromStrict . TE.encodeUtf8 + +memoryRowColumns :: Text +memoryRowColumns = + "memory_id, agent_id, session_id, namespace, scope_kind, scope_ref, memory_type, content, priority, confidence, tags::text, status, superseded_by, supersedes, created_at, updated_at" + +qualifiedMemoryRowColumns :: Text -> Text +qualifiedMemoryRowColumns prefix = + prefix + <> ".memory_id, " + <> prefix + <> ".agent_id, " + <> prefix + <> ".session_id, " + <> prefix + <> ".namespace, " + <> prefix + <> ".scope_kind, " + <> prefix + <> ".scope_ref, " + <> prefix + <> ".memory_type, " + <> prefix + <> ".content, " + <> prefix + <> ".priority, " + <> prefix + <> ".confidence, " + <> prefix + <> ".tags::text, " + <> prefix + <> ".status, " + <> prefix + <> ".superseded_by, " + <> prefix + <> ".supersedes, " + <> prefix + <> ".created_at, " + <> prefix + <> ".updated_at" + +selectMemoryByIdStmt :: Statement Text (Maybe MemoryRow) +selectMemoryByIdStmt = + preparable + ( "SELECT " + <> memoryRowColumns + <> " FROM kioku_memories WHERE memory_id = $1" + ) + (E.param (E.nonNullable E.text)) + (D.rowMaybe memoryRowDecoder) + +selectActiveByNamespaceStmt :: Statement Text [MemoryRecord] +selectActiveByNamespaceStmt = + preparable + ( "SELECT " + <> memoryRowColumns + <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 ORDER BY priority ASC, created_at DESC" + ) + (E.param (E.nonNullable E.text)) + (D.rowList memoryRecordDecoder) + +selectActiveByNamespaceRowsStmt :: Statement Text [MemoryRow] +selectActiveByNamespaceRowsStmt = + preparable + ( "SELECT " + <> memoryRowColumns + <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 ORDER BY priority ASC, created_at DESC" + ) + (E.param (E.nonNullable E.text)) + (D.rowList memoryRowDecoder) + +-- | Active memories carrying __exactly__ the given scope. +-- +-- Recall searches namespace-wide for a global scope; scoped reads are exact-scope. The +-- predicate below /requires/ @scope_kind@ and @scope_ref@ to be NULL for a global scope, +-- where 'Kioku.Recall.selectFtsCandidatesStmt' would drop the scope filter entirely and +-- return the whole namespace. Both behaviours are intentional; see docs/user/recall.md. +selectActiveByScopeStmt :: Statement (Text, Maybe Text, Maybe Text) [MemoryRecord] +selectActiveByScopeStmt = + preparable + ( "SELECT " + <> memoryRowColumns + <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 AND ((scope_kind = $2 AND scope_ref = $3) OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) ORDER BY priority ASC, created_at DESC" + ) + ( contrazip3 + (E.param (E.nonNullable E.text)) + (E.param (E.nullable E.text)) + (E.param (E.nullable E.text)) + ) + (D.rowList memoryRecordDecoder) + +selectActiveByScopeRowsStmt :: Statement (Text, Maybe Text, Maybe Text) [MemoryRow] +selectActiveByScopeRowsStmt = + preparable + ( "SELECT " + <> memoryRowColumns + <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 AND ((scope_kind = $2 AND scope_ref = $3) OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) ORDER BY priority ASC, created_at DESC" + ) + ( contrazip3 + (E.param (E.nonNullable E.text)) + (E.param (E.nullable E.text)) + (E.param (E.nullable E.text)) + ) + (D.rowList memoryRowDecoder) + +selectBySessionStmt :: Statement Text [MemoryRecord] +selectBySessionStmt = + preparable + ( "SELECT " + <> memoryRowColumns + <> " FROM kioku_memories WHERE session_id = $1 ORDER BY created_at DESC" + ) + (E.param (E.nonNullable E.text)) + (D.rowList memoryRecordDecoder) + +selectBySessionRowsStmt :: Statement Text [MemoryRow] +selectBySessionRowsStmt = + preparable + ( "SELECT " + <> memoryRowColumns + <> " FROM kioku_memories WHERE session_id = $1 ORDER BY created_at DESC" + ) + (E.param (E.nonNullable E.text)) + (D.rowList memoryRowDecoder) + +selectByTypeStmt :: Statement (Text, Text) [MemoryRecord] +selectByTypeStmt = + preparable + ( "SELECT " + <> memoryRowColumns + <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 AND memory_type = $2 ORDER BY priority ASC, created_at DESC" + ) + ( contrazip2 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.text)) + ) + (D.rowList memoryRecordDecoder) + +selectByTypeRowsStmt :: Statement (Text, Text) [MemoryRow] +selectByTypeRowsStmt = + preparable + ( "SELECT " + <> memoryRowColumns + <> " FROM kioku_memories WHERE status = 'active' AND namespace = $1 AND memory_type = $2 ORDER BY priority ASC, created_at DESC" + ) + ( contrazip2 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.text)) + ) + (D.rowList memoryRowDecoder) + +selectSupersessionChainStmt :: Statement Text [MemoryRow] +selectSupersessionChainStmt = + preparable + ( "WITH RECURSIVE chain AS (" + <> "SELECT " + <> memoryRowColumns + <> " FROM kioku_memories WHERE memory_id = $1 " + <> "UNION " + <> "SELECT " + <> qualifiedMemoryRowColumns "m" + <> " FROM kioku_memories m JOIN chain c ON " + <> "m.memory_id = c.supersedes " + <> "OR m.supersedes = c.memory_id " + <> "OR m.memory_id = c.superseded_by " + <> "OR m.superseded_by = c.memory_id" + <> ") SELECT " + <> memoryRowColumns + <> " FROM chain ORDER BY created_at ASC, memory_id ASC" + ) + (E.param (E.nonNullable E.text)) + (D.rowList memoryRowDecoder) + +upsertMemoryStmt :: Statement MemoryRow () +upsertMemoryStmt = + preparable + """ + INSERT INTO kioku_memories + (memory_id, agent_id, session_id, namespace, scope_kind, scope_ref, memory_type, content, + priority, confidence, tags, status, superseded_by, supersedes, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11::jsonb, $12, $13, $14, $15, $16) + ON CONFLICT (memory_id) DO UPDATE SET + agent_id = EXCLUDED.agent_id, + session_id = EXCLUDED.session_id, + namespace = EXCLUDED.namespace, + scope_kind = EXCLUDED.scope_kind, + scope_ref = EXCLUDED.scope_ref, + memory_type = EXCLUDED.memory_type, + content = EXCLUDED.content, + priority = EXCLUDED.priority, + confidence = EXCLUDED.confidence, + tags = EXCLUDED.tags, + status = EXCLUDED.status, + superseded_by = EXCLUDED.superseded_by, + supersedes = EXCLUDED.supersedes, + updated_at = EXCLUDED.updated_at + """ + memoryRowEncoder + D.noResult + +memoryRowEncoder :: E.Params MemoryRow +memoryRowEncoder = + ((\row -> row.memoryId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.agentId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.sessionId) >$< E.param (E.nullable E.text)) + <> ((\row -> row.namespace) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.scopeKind) >$< E.param (E.nullable E.text)) + <> ((\row -> row.scopeRef) >$< E.param (E.nullable E.text)) + <> ((\row -> row.memoryType) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.content) >$< E.param (E.nonNullable E.text)) + <> ((fromIntegral @Int @Int32 . \row -> row.priority) >$< E.param (E.nonNullable E.int4)) + <> ((\row -> row.confidence) >$< E.param (E.nonNullable E.text)) + <> ((encodeTags . (\row -> row.tags)) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.status) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.supersededBy) >$< E.param (E.nullable E.text)) + <> ((\row -> row.supersedes) >$< E.param (E.nullable E.text)) + <> ((\row -> row.createdAt) >$< E.param (E.nonNullable E.timestamptz)) + <> ((\row -> row.updatedAt) >$< E.param (E.nonNullable E.timestamptz)) + +updateMemorySupersededStmt :: Statement (Text, Text, UTCTime) () +updateMemorySupersededStmt = + preparable + "UPDATE kioku_memories SET status = 'superseded', superseded_by = $2, updated_at = $3 WHERE memory_id = $1" + ( contrazip3 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.timestamptz)) + ) + D.noResult + +updateMemoryArchivedStmt :: Statement (Text, UTCTime) () +updateMemoryArchivedStmt = + preparable + "UPDATE kioku_memories SET status = 'archived', updated_at = $2 WHERE memory_id = $1" + (contrazip2 (E.param (E.nonNullable E.text)) (E.param (E.nonNullable E.timestamptz))) + D.noResult + +updateMemoryTagsStmt :: Statement (Text, Set Text, UTCTime) () +updateMemoryTagsStmt = + preparable + "UPDATE kioku_memories SET tags = $2::jsonb, updated_at = $3 WHERE memory_id = $1" + ( contrazip3 + (E.param (E.nonNullable E.text)) + (encodeTags >$< E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.timestamptz)) + ) + D.noResult + +updateMemoryConfidenceStmt :: Statement (Text, Text, UTCTime) () +updateMemoryConfidenceStmt = + preparable + "UPDATE kioku_memories SET confidence = $2, updated_at = $3 WHERE memory_id = $1" + ( contrazip3 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.timestamptz)) + ) + D.noResult + +updateMemoryMergedStmt :: Statement (Text, Text, UTCTime) () +updateMemoryMergedStmt = + preparable + "UPDATE kioku_memories SET status = 'merged', superseded_by = $2, updated_at = $3 WHERE memory_id = $1" + ( contrazip3 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.timestamptz)) + ) + D.noResult diff --git a/src/Kioku/ReadModel.hs b/src/Kioku/ReadModel.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/ReadModel.hs @@ -0,0 +1,156 @@ +-- | The schema identities of every Kioku read model registered in the +-- @keiro_read_models@ registry. +-- +-- Kioku evolves its projections with /additive/ migrations: new columns are +-- added and backfill with sensible defaults, so existing rows remain correct +-- for the newer read-model version without a full offline rebuild. When such a +-- migration ships, the code-declared 'Keiro.ReadModel.version' / +-- 'Keiro.ReadModel.shapeHash' advance, but any pre-existing +-- @keiro_read_models@ row stays pinned at the old identity — +-- 'Keiro.ReadModel.registerReadModel' only inserts, it never bumps an existing +-- row. Every subsequent query then fails closed with +-- 'Keiro.ReadModel.ReadModelStaleSchema'. +-- +-- 'reconcileReadModelRegistry' repairs those rows to the identity the compiled +-- code expects, deriving every name, version, and shape hash from +-- 'kiokuReadModelSchemas' — the same 'ReadModel' values the queries use, so the +-- registry can never disagree with the code. The @kioku-migrate@ executable runs +-- it immediately after applying migrations, which is why a read-model version +-- bump needs no hand-written registry SQL. A host that applies migrations as a +-- library (by running @Kioku.Migrations.kiokuMigrationPlan@ through pg-migrate) +-- must call it itself; see @docs\/user\/library-api.md@. +-- +-- Because Kioku guarantees its read-model migrations leave the table data correct +-- for the current version, advancing the registry guard is safe here without a +-- rebuild. A consumer that introduces a /non-additive/ reshape must instead +-- rewrite the projected data in its migration before reconciling. +module Kioku.ReadModel + ( ReadModelSchema (..), + kiokuReadModelSchemas, + registerKiokuReadModels, + ReconcileOutcome (..), + reconcileReadModelRegistry, + ) +where + +import Effectful (Eff, (:>)) +import Keiro.ReadModel (ReadModel (..)) +import Keiro.ReadModel.Schema qualified as Schema +import Kioku.Memory.ReadModel + ( memoriesByNamespaceReadModel, + memoriesByNamespaceRowsReadModel, + memoriesByScopeReadModel, + memoriesByScopeRowsReadModel, + memoriesBySessionReadModel, + memoriesBySessionRowsReadModel, + memoriesByTypeReadModel, + memoriesByTypeRowsReadModel, + memoryByIdReadModel, + memorySupersessionChainReadModel, + ) +import Kioku.Prelude +import Kioku.Session.ReadModel + ( awaitingSessionsByCorrelationKeyReadModel, + sessionByIdReadModel, + sessionChainReadModel, + sessionDelegationChildrenReadModel, + sessionsByFocusReadModel, + sessionsByNamespaceReadModel, + sessionsByScopeReadModel, + sessionsByStartedRangeReadModel, + turnsBySessionReadModel, + ) +import Kiroku.Store.Effect (Store) + +-- | The registry identity of a read model: its logical name plus the schema +-- 'version' and 'shapeHash' the current code expects. +data ReadModelSchema = ReadModelSchema + { readModelName :: !Text, + readModelVersion :: !Int, + readModelShapeHash :: !Text + } + deriving stock (Eq, Show) + +schemaOf :: ReadModel q r -> ReadModelSchema +schemaOf rm = ReadModelSchema rm.name rm.version rm.shapeHash + +-- | Every Kioku read model paired with the schema identity the current code +-- expects. Ordered session models first, then memory models. +kiokuReadModelSchemas :: [ReadModelSchema] +kiokuReadModelSchemas = + [ schemaOf sessionByIdReadModel, + schemaOf sessionsByNamespaceReadModel, + schemaOf sessionsByScopeReadModel, + schemaOf sessionsByFocusReadModel, + schemaOf sessionsByStartedRangeReadModel, + schemaOf sessionChainReadModel, + schemaOf sessionDelegationChildrenReadModel, + schemaOf awaitingSessionsByCorrelationKeyReadModel, + schemaOf turnsBySessionReadModel, + schemaOf memoryByIdReadModel, + schemaOf memoriesByNamespaceReadModel, + schemaOf memoriesByNamespaceRowsReadModel, + schemaOf memoriesByScopeReadModel, + schemaOf memoriesByScopeRowsReadModel, + schemaOf memoriesBySessionReadModel, + schemaOf memoriesBySessionRowsReadModel, + schemaOf memoriesByTypeReadModel, + schemaOf memoriesByTypeRowsReadModel, + schemaOf memorySupersessionChainReadModel + ] + +-- | Register every Kioku read model at application startup. +-- +-- Keiro 0.3 deliberately stopped registering models on their first query. This +-- operation is idempotent and leaves an existing row unchanged, allowing Keiro +-- to continue failing closed when its version or shape hash is stale. +registerKiokuReadModels :: (Store :> es) => Eff es () +registerKiokuReadModels = + forM_ kiokuReadModelSchemas \schema -> + void $ + Schema.registerReadModel + schema.readModelName + schema.readModelVersion + schema.readModelShapeHash + +-- | What reconciliation did to one read model's registry row. +data ReconcileOutcome + = -- | No row existed; one was inserted at the current identity. + Registered + | -- | The row disagreed with the code and was bumped to the current identity. + Reconciled + | -- | The row already matched the code. Left untouched, status and all. + AlreadyCurrent + deriving stock (Eq, Show) + +-- | Bring every row of keiro's @keiro_read_models@ registry up to the schema +-- identity the compiled code expects, and report what changed. +-- +-- Idempotent: a second run reports 'AlreadyCurrent' for every model and writes +-- nothing. It is built on keiro's own registry statements, which are compiled +-- into whichever keiro version Kioku links against — so it finds the registry +-- table wherever that version puts it, and survives keiro's schema relocation +-- without change. +reconcileReadModelRegistry :: + (Store :> es) => Eff es [(ReadModelSchema, ReconcileOutcome)] +reconcileReadModelRegistry = + forM kiokuReadModelSchemas \schema -> do + existing <- Schema.lookupReadModel schema.readModelName + outcome <- case existing of + Nothing -> + Registered + <$ Schema.registerReadModel + schema.readModelName + schema.readModelVersion + schema.readModelShapeHash + Just metadata + | metadata.version == schema.readModelVersion + && metadata.shapeHash == schema.readModelShapeHash -> + pure AlreadyCurrent + | otherwise -> + Reconciled + <$ Schema.markLive + schema.readModelName + schema.readModelVersion + schema.readModelShapeHash + pure (schema, outcome) diff --git a/src/Kioku/Recall.hs b/src/Kioku/Recall.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Recall.hs @@ -0,0 +1,838 @@ +module Kioku.Recall + ( RecallStrategy (..), + RecallRequest (..), + RecallHit (..), + RecallExecutionPlan (..), + recall, + planRecallExecution, + fuseRecallCandidates, + blendScore, + rrfTerm, + recencyDecay, + priorityWeight, + confidenceWeight, + applyCharacterBudgets, + getActiveInNamespace, + getActiveByScope, + getGlobal, + getById, + getBySession, + getByType, + + -- * Test seams + -- $testSeams + selectFtsCandidates, + selectVectorCandidates, + vectorLiteral, + selectVectorCandidatesStmt, + selectVectorCandidatesExactStmt, + selectVectorCandidatesDiagnosed, + VectorChannelOutcome (..), + vectorChannelStarved, + VectorCandidateQuery, + vectorCandidateQuery, + memoryRecordColumns, + candidatePoolSize, + ) +where + +import Baikai.Embedding (EmbeddingModel) +import Data.Aeson qualified as Aeson +import Data.ByteString (ByteString) +import Data.ByteString.Lazy qualified as BL +import Data.Functor.Contravariant ((>$<)) +import Data.Int (Int32) +import Data.List qualified as List +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Ord (Down (..)) +import Data.Set qualified as Set +import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE +import Data.Time (diffUTCTime) +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import Effectful (Eff, IOE, (:>)) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Keiro.ReadModel (ConsistencyMode (..), ReadModelError, runQueryWith) +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), scopeFromColumns, scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Api.Types (MemoryRecord (..), MemoryType, memoryTypeToText) +import Kioku.Id (MemoryId, SessionId, idText) +import Kioku.Memory.Embedding (embedWithRetry) +import Kioku.Memory.ReadModel + ( MemoriesByNamespaceQuery (..), + MemoriesByScopeQuery (..), + MemoriesBySessionQuery (..), + MemoriesByTypeQuery (..), + MemoryByIdQuery (..), + MemoryRow (..), + memoriesByNamespaceReadModel, + memoriesByScopeReadModel, + memoriesBySessionReadModel, + memoriesByTypeReadModel, + memoryByIdReadModel, + ) +import Kioku.Prelude +import Kioku.Recall.Capability (VectorCapability (..)) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Transaction (runTransaction) + +-- $testSeams +-- Exported so the candidate SQL can be exercised directly against a real database +-- (@Kioku.RecallSqlSpec@) rather than only through 'recall', which would drag in an +-- embedding endpoint. They are not part of the intended public API. +-- +-- 'selectVectorCandidatesStmt', 'vectorCandidateQuery', 'memoryRecordColumns' and +-- 'candidatePoolSize' are exported for @Kioku.RecallHarness@, the recall-quality instrument. It +-- needs the statement itself (rather than 'selectVectorCandidates', which wraps it in its own +-- transaction) so that it can run it under a @SET LOCAL@, and it needs the projection and the +-- pool size so that its @EXPLAIN@ describes the query that actually runs. That last one is not a +-- nicety: the projection's row width sets the cost of the top-N sort the /exact/ plan needs, and +-- that cost is what the planner weighs against the HNSW scan — so an @EXPLAIN@ carrying a +-- different select list can silently choose a different plan and report a different answer. +-- Restating them in the harness rather than exporting them is how the harness got that wrong +-- once. + +data RecallStrategy = Keyword | Embedding | Hybrid + deriving stock (Generic, Eq, Show) + +-- | A recall request. +-- +-- __Global scope means "namespace-wide" here.__ Recall searches namespace-wide for a global +-- scope; scoped reads are exact-scope. A 'ScopeGlobal' request returns every active memory in +-- the namespace, entity-scoped rows included — the scope filter simply vanishes. That is the +-- opposite of what 'getActiveByScope' does with the same value. See docs/user/recall.md. +data RecallRequest = RecallRequest + { -- | 'ScopeGlobal' searches the whole namespace; an entity scope matches exactly. + scope :: !MemoryScope, + query :: !Text, + strategy :: !RecallStrategy, + maxResults :: !Int + } + deriving stock (Generic, Eq, Show) + +data RecallHit = RecallHit + { memory :: !MemoryRecord, + score :: !Double, + ftsRank :: !(Maybe Int), + vecRank :: !(Maybe Int) + } + deriving stock (Generic, Eq, Show) + +data RecallExecutionPlan = RecallExecutionPlan + { runFts :: !Bool, + runVector :: !Bool, + needsQueryEmbedding :: !Bool + } + deriving stock (Generic, Eq, Show) + +data RecallCandidateQuery = RecallCandidateQuery + { query :: !Text, + namespace :: !Text, + scopeKind :: !(Maybe Text), + scopeRef :: !(Maybe Text), + limit :: !Int32 + } + deriving stock (Generic, Eq, Show) + +data VectorCandidateQuery = VectorCandidateQuery + { queryVector :: !Text, + namespace :: !Text, + scopeKind :: !(Maybe Text), + scopeRef :: !(Maybe Text), + limit :: !Int32 + } + deriving stock (Generic, Eq, Show) + +data FusedCandidate = FusedCandidate + { memory :: !MemoryRecord, + ftsRank :: !(Maybe Int), + vecRank :: !(Maybe Int) + } + deriving stock (Generic, Eq, Show) + +-- | Run a recall request: plan, optionally embed the query, select candidates from each +-- active channel, fuse by reciprocal rank, score, and trim. +-- +-- Recall searches namespace-wide for a global scope; scoped reads are exact-scope. +recall :: + (IOE :> es, Store :> es) => + EmbeddingModel -> + VectorCapability -> + RecallRequest -> + Eff es [RecallHit] +recall model capability req = do + now <- liftIO getCurrentTime + executeRecallPlan now model req (planRecallExecution capability req.strategy) + +planRecallExecution :: VectorCapability -> RecallStrategy -> RecallExecutionPlan +planRecallExecution capability strategy = + case capability of + VectorAvailable -> + case strategy of + Keyword -> RecallExecutionPlan {runFts = True, runVector = False, needsQueryEmbedding = False} + Embedding -> RecallExecutionPlan {runFts = False, runVector = True, needsQueryEmbedding = True} + Hybrid -> RecallExecutionPlan {runFts = True, runVector = True, needsQueryEmbedding = True} + VectorExtensionUnavailable -> + keywordExecutionPlan + VectorColumnsUnavailable _ -> + keywordExecutionPlan + -- A dimension mismatch is a configuration error, not a missing feature, but recall's + -- response is the same: the vector channel cannot work, so degrade to keyword rather + -- than fail. The worker is where it is reported loudly. + VectorDimensionMismatch {} -> + keywordExecutionPlan + +keywordExecutionPlan :: RecallExecutionPlan +keywordExecutionPlan = + RecallExecutionPlan {runFts = True, runVector = False, needsQueryEmbedding = False} + +executeRecallPlan :: + (IOE :> es, Store :> es) => + UTCTime -> + EmbeddingModel -> + RecallRequest -> + RecallExecutionPlan -> + Eff es [RecallHit] +executeRecallPlan now model req execution + | execution.needsQueryEmbedding = + embedThenRecall now model req execution + | otherwise = do + ftsRows <- selectIf execution.runFts (selectFtsCandidates req) + pure (finishRecall now req ftsRows []) + +embedThenRecall :: + (IOE :> es, Store :> es) => + UTCTime -> + EmbeddingModel -> + RecallRequest -> + RecallExecutionPlan -> + Eff es [RecallHit] +embedThenRecall now model req execution = do + embedded <- liftIO (embedWithRetry model 2 req.query) + case embedded of + Left _err -> + keywordOnly now req + Right queryVector -> do + ftsRows <- selectIf execution.runFts (selectFtsCandidates req) + vecRows <- selectIf execution.runVector (selectVectorCandidates req queryVector) + pure (finishRecall now req ftsRows vecRows) + +keywordOnly :: + (Store :> es) => + UTCTime -> + RecallRequest -> + Eff es [RecallHit] +keywordOnly now req = do + ftsRows <- selectFtsCandidates req + pure (finishRecall now req ftsRows []) + +selectIf :: (Applicative f) => Bool -> f [a] -> f [a] +selectIf True action = action +selectIf False _ = pure [] + +finishRecall :: UTCTime -> RecallRequest -> [MemoryRecord] -> [MemoryRecord] -> [RecallHit] +finishRecall now req ftsRows vecRows = + applyCharacterBudgets perMemoryCharacterBudget totalCharacterBudget $ + take (max 0 req.maxResults) $ + fuseRecallCandidates now ftsRows vecRows + +selectFtsCandidates :: + (Store :> es) => + RecallRequest -> + Eff es [MemoryRecord] +selectFtsCandidates req = + runTransaction $ + Tx.statement + (candidateQuery req) + selectFtsCandidatesStmt + +-- | What the vector channel did, so that a degraded semantic half stops being invisible. +-- +-- This exists because of how the defect it describes survived. 'fuseRecallCandidates' blends the +-- two channels by rank, so a vector channel that returns nothing contributes no ranks and the +-- score decays smoothly into pure keyword scoring: no error, no warning, and nothing in a +-- 'RecallHit' recording that the semantic half of a "hybrid" search came back empty. A caller +-- looking at plausible keyword results had no way to know. +data VectorChannelOutcome = VectorChannelOutcome + { -- | Rows the approximate (HNSW) pass returned. + annRows :: !Int, + -- | Whether the exact pass ran because the approximate one came back short of the pool. + exactFallbackFired :: !Bool, + -- | Rows finally handed to the fusion. + rowsReturned :: !Int + } + deriving stock (Generic, Eq, Show) + +-- | Did the approximate pass miss rows that were really there? +-- +-- True exactly when the exact fallback found more than the ANN pass did — that is, when the ANN +-- scan starved and the fallback rescued it. A host that wants a metric or a log line for the +-- health of its semantic channel should count these. +vectorChannelStarved :: VectorChannelOutcome -> Bool +vectorChannelStarved outcome = + outcome.exactFallbackFired && outcome.rowsReturned > outcome.annRows + +selectVectorCandidates :: + (Store :> es) => + RecallRequest -> + Vector Double -> + Eff es [MemoryRecord] +selectVectorCandidates req queryVector = + snd <$> selectVectorCandidatesDiagnosed req queryVector + +-- | The vector channel, with a report of what it did. See 'VectorChannelOutcome'. +-- +-- == Why there are two passes +-- +-- The HNSW index is /post-filtered/: it picks candidates by distance alone, and the namespace, +-- scope, and @status@ predicates are applied afterwards, to rows it has already chosen. When the +-- memories nearest the query sit outside the caller's scope — a small scope inside a large +-- namespace, which is the normal shape of kioku data — the index spends its whole budget on rows +-- the filter then discards and the channel returns __nothing__. +-- +-- So: run the approximate pass, and if it comes back short of the pool, run an exact pass that +-- cannot starve. The trigger is "short of the pool" rather than "empty" because a partial +-- starvation is a starvation too, and because the exact pass is cheap in precisely the case that +-- makes it fire spuriously — a scope holding fewer than 'candidatePoolSize' embedded memories, +-- where scanning all of them costs nothing. +-- +-- == Why the exact pass rather than pgvector's own remedy +-- +-- pgvector 0.8's @hnsw.iterative_scan@ is designed for exactly this, and it was measured across +-- five freshly built indexes on a 20000-row starving corpus. It returned the right answer 2 times +-- in 5 (@relaxed_order@) and 4 times in 5 (@strict_order@). HNSW construction is randomized, so +-- whether the iterative scan reaches the in-scope rows within its budget depends on the graph it +-- happened to get. A remedy that works 40% of the time is not a remedy — and, worse, it passes +-- any single test run. The exact pass returned the right answer 5 times in 5, needs no minimum +-- pgvector version, and cannot starve by construction. +-- +-- == The cost, honestly +-- +-- The exact pass scans every embedded row in the caller's scope: about 7ms per 2000 rows on the +-- measurement machine, growing linearly. That cost is paid only when the approximate pass came +-- back short, and it is bounded by the size of the scope the caller asked about — which is the +-- set they wanted searched anyway. On a large scope with no selective filter the approximate pass +-- fills the pool and the fallback never fires, so the common path is unchanged. +selectVectorCandidatesDiagnosed :: + (Store :> es) => + RecallRequest -> + Vector Double -> + Eff es (VectorChannelOutcome, [MemoryRecord]) +selectVectorCandidatesDiagnosed req queryVector = + runTransaction do + -- The HNSW scan visits at most @hnsw.ef_search@ candidates, and the default (40) is below the + -- pool size (50) — so even with nothing for the filter to discard, the pool never fills and + -- the channel silently under-delivers by 20%. Measured: 40 rows returned for a LIMIT of 50 on + -- a 2000-row corpus with no out-of-scope rows at all. + -- + -- (The comment that used to live at 'candidatePoolSize' claimed pgvector searches with + -- @ef = max(ef_search, LIMIT)@ so the pool fills at the default. On pgvector 0.8.2 it does + -- not. That claim was inherited, plausible, and false.) + -- + -- Raising it to exactly the pool size fills the pool at no measurable cost (0.149ms vs + -- 0.138ms) and does not move the planner off the HNSW path. Raising it far higher — the + -- remedy a previous plan prescribed — /does/ move the planner, onto an ANN scan that then + -- starves, which is how this defect was originally mis-diagnosed. + Tx.sql efSearchSetting + annRows <- Tx.statement query selectVectorCandidatesStmt + if length annRows >= fromIntegral candidatePoolSize + then + pure + ( VectorChannelOutcome + { annRows = length annRows, + exactFallbackFired = False, + rowsReturned = length annRows + }, + annRows + ) + else do + exactRows <- Tx.statement query selectVectorCandidatesExactStmt + pure + ( VectorChannelOutcome + { annRows = length annRows, + exactFallbackFired = True, + rowsReturned = length exactRows + }, + exactRows + ) + where + query = vectorCandidateQuery req queryVector + +-- | @SET LOCAL@, so it lives exactly as long as the transaction the query runs in and cannot +-- leak into the rest of the connection. +efSearchSetting :: ByteString +efSearchSetting = + TE.encodeUtf8 ("SET LOCAL hnsw.ef_search = " <> Text.pack (show candidatePoolSize)) + +candidateQuery :: RecallRequest -> RecallCandidateQuery +candidateQuery req = + RecallCandidateQuery + { query = req.query, + namespace = scopeNamespaceText req.scope, + scopeKind = scopeKindText req.scope, + scopeRef = scopeRefText req.scope, + limit = candidatePoolSize + } + +vectorCandidateQuery :: RecallRequest -> Vector Double -> VectorCandidateQuery +vectorCandidateQuery req queryVector = + VectorCandidateQuery + { queryVector = vectorLiteral queryVector, + namespace = scopeNamespaceText req.scope, + scopeKind = scopeKindText req.scope, + scopeRef = scopeRefText req.scope, + limit = candidatePoolSize + } + +fuseRecallCandidates :: UTCTime -> [MemoryRecord] -> [MemoryRecord] -> [RecallHit] +fuseRecallCandidates now ftsRows vecRows = + List.sortOn (Down . (\hit -> hit.score)) $ + toHit <$> Map.elems fused + where + fused = + foldRanked (\rank row -> upsertFts row rank) Map.empty ftsRows + & \m -> foldRanked (\rank row -> upsertVec row rank) m vecRows + + toHit candidate = + RecallHit + { memory = candidate.memory, + score = blendScore now candidate.memory candidate.ftsRank candidate.vecRank, + ftsRank = candidate.ftsRank, + vecRank = candidate.vecRank + } + +foldRanked :: (Int -> MemoryRecord -> Map Text FusedCandidate -> Map Text FusedCandidate) -> Map Text FusedCandidate -> [MemoryRecord] -> Map Text FusedCandidate +foldRanked f initial rows = + foldl + (\acc (rank, row) -> f rank row acc) + initial + (zip [1 ..] rows) + +upsertFts :: MemoryRecord -> Int -> Map Text FusedCandidate -> Map Text FusedCandidate +upsertFts row rank = + Map.alter (Just . addRank) row.memoryId + where + addRank :: Maybe FusedCandidate -> FusedCandidate + addRank Nothing = + FusedCandidate {memory = row, ftsRank = Just rank, vecRank = Nothing} + addRank (Just existing) = + FusedCandidate + { memory = existing.memory, + ftsRank = existing.ftsRank <|> Just rank, + vecRank = existing.vecRank + } + +upsertVec :: MemoryRecord -> Int -> Map Text FusedCandidate -> Map Text FusedCandidate +upsertVec row rank = + Map.alter (Just . addRank) row.memoryId + where + addRank :: Maybe FusedCandidate -> FusedCandidate + addRank Nothing = + FusedCandidate {memory = row, ftsRank = Nothing, vecRank = Just rank} + addRank (Just existing) = + FusedCandidate + { memory = existing.memory, + ftsRank = existing.ftsRank, + vecRank = existing.vecRank <|> Just rank + } + +blendScore :: UTCTime -> MemoryRecord -> Maybe Int -> Maybe Int -> Double +blendScore now memory ftsRank vecRank = + maybe 0 rrfTerm ftsRank + + maybe 0 rrfTerm vecRank + + recencyWeight * recencyDecay now memory.createdAt + + prioritySignalWeight * priorityWeight memory.priority + + confidenceSignalWeight * confidenceWeight memory.confidence + +rrfTerm :: Int -> Double +rrfTerm rank = + 1 / (rrfK + fromIntegral rank) + +recencyDecay :: UTCTime -> UTCTime -> Double +recencyDecay now createdAt = + exp (negate (log 2) * ageDays / recencyHalfLifeDays) + where + ageDays = max 0 (realToFrac (diffUTCTime now createdAt) / secondsPerDay) + +priorityWeight :: Int -> Double +priorityWeight priority + | priority <= alwaysInjectPriority = 1 + | otherwise = clamp01 (1 - (fromIntegral priority / priorityMax)) + +confidenceWeight :: Text -> Double +confidenceWeight = \case + "high" -> 1 + "medium" -> 0.6 + "low" -> 0.3 + _ -> 0.3 + +applyCharacterBudgets :: Int -> Int -> [RecallHit] -> [RecallHit] +applyCharacterBudgets perMemoryCap totalCap = + go 0 [] + where + go _ acc [] = reverse acc + go used acc (hit : rest) + | totalCap <= 0 = reverse acc + | used + Text.length truncated.memory.content > totalCap = reverse acc + | otherwise = go (used + Text.length truncated.memory.content) (truncated : acc) rest + where + truncated = truncateHit perMemoryCap hit + +truncateHit :: Int -> RecallHit -> RecallHit +truncateHit cap hit = + RecallHit + { memory = truncateMemory cap hit.memory, + score = hit.score, + ftsRank = hit.ftsRank, + vecRank = hit.vecRank + } + +truncateMemory :: Int -> MemoryRecord -> MemoryRecord +truncateMemory cap row = + MemoryRecord + { memoryId = row.memoryId, + agentId = row.agentId, + sessionId = row.sessionId, + scope = row.scope, + memoryType = row.memoryType, + content = truncateText cap row.content, + priority = row.priority, + confidence = row.confidence, + tags = row.tags, + status = row.status, + createdAt = row.createdAt + } + +truncateText :: Int -> Text -> Text +truncateText cap content + | cap <= 0 = "" + | Text.length content <= cap = content + | cap <= Text.length ellipsis = Text.take cap ellipsis + | otherwise = Text.take (cap - Text.length ellipsis) content <> ellipsis + +clamp01 :: Double -> Double +clamp01 = max 0 . min 1 + +-- | Full-text candidates. +-- +-- The scope predicate @(($3 IS NULL AND $4 IS NULL) OR (scope_kind = $3 AND scope_ref = $4))@ +-- is why recall searches namespace-wide for a global scope; scoped reads are exact-scope. For +-- a global scope both parameters are NULL, the first disjunct is always true, and the filter +-- vanishes. 'Kioku.Memory.ReadModel.selectActiveByScopeStmt' requires the columns to be NULL +-- instead. The @ORDER BY@ here is free to carry a recency tiebreak: a GIN index provides no +-- ordering, so there is no pathkey to preserve. +selectFtsCandidatesStmt :: Statement RecallCandidateQuery [MemoryRecord] +selectFtsCandidatesStmt = + preparable + ( "SELECT " + <> memoryRecordColumns + <> """ + FROM kiroku.kioku_memories + WHERE status = 'active' + AND namespace = $2 + AND (($3 IS NULL AND $4 IS NULL) OR (scope_kind = $3 AND scope_ref = $4)) + AND content_tsv @@ websearch_to_tsquery('english', $1) + ORDER BY ts_rank(content_tsv, websearch_to_tsquery('english', $1)) DESC, created_at DESC + LIMIT $5 + """ + ) + recallCandidateQueryEncoder + (D.rowList memoryRecordDecoder) + +-- | Vector candidates, ordered by cosine distance and nothing else. +-- +-- The @ORDER BY@ is deliberately a single expression. An HNSW index can only produce the +-- distance pathkey, so a second sort key (this query used to carry @created_at DESC@) leaves +-- the planner to make up the difference: on PostgreSQL 13+ it bolts an @Incremental Sort@ on +-- top of the index scan, and where incremental sort is unavailable or disabled it abandons +-- the index entirely for a sequential scan plus a full sort. Ordering by distance alone makes +-- the index scan unconditional. Nothing is lost: the statement does not return the distance, +-- so a caller could not re-break ties anyway, and exact ties between 1536-dimension float +-- vectors essentially do not occur. +-- +-- Recall that this is a *post-filtered* ANN scan: the namespace, scope and status predicates +-- are applied to rows the index has already chosen by distance. See 'candidatePoolSize' for +-- what that costs. +-- +-- The scope predicate is the same one 'selectFtsCandidatesStmt' carries: recall searches +-- namespace-wide for a global scope; scoped reads are exact-scope. +selectVectorCandidatesStmt :: Statement VectorCandidateQuery [MemoryRecord] +selectVectorCandidatesStmt = + preparable + ( "SELECT " + <> memoryRecordColumns + <> """ + FROM kiroku.kioku_memories + WHERE status = 'active' + AND namespace = $2 + AND (($3 IS NULL AND $4 IS NULL) OR (scope_kind = $3 AND scope_ref = $4)) + AND embedding IS NOT NULL + ORDER BY embedding <=> $1::vector + LIMIT $5 + """ + ) + vectorCandidateQueryEncoder + (D.rowList memoryRecordDecoder) + +-- | The exact vector scan: every embedded row in the caller's scope, ranked by distance, top-N. +-- It cannot starve, because the filter is applied /before/ the ranking rather than after it. +-- +-- The @OFFSET 0@ is the whole mechanism and must not be "tidied away". It is an optimisation +-- fence: it stops Postgres from pulling the subquery up into the outer query, which in turn stops +-- the outer @ORDER BY embedding <=> …@ from reaching the HNSW index. Without it the planner +-- flattens the two levels back into 'selectVectorCandidatesStmt' and we are measuring — and +-- shipping — the very query we are trying to avoid. +-- +-- A @MATERIALIZED@ CTE would also fence it, and was rejected: materialising forces every in-scope +-- row's 1536-dimension embedding into memory (about 6KB each, so ~120MB for a 20000-row scope), +-- whereas the fence streams and the top-N sort holds only 50 rows. +-- +-- The predicates are identical to 'selectVectorCandidatesStmt''s, including @embedding IS NOT +-- NULL@ — here it is a correctness filter rather than an index-matching one, but it must stay +-- either way, since a NULL embedding has no distance to anything. +selectVectorCandidatesExactStmt :: Statement VectorCandidateQuery [MemoryRecord] +selectVectorCandidatesExactStmt = + preparable + ( "SELECT " + <> memoryRecordColumns + <> """ + FROM (SELECT * + FROM kiroku.kioku_memories + WHERE status = 'active' + AND namespace = $2 + AND (($3 IS NULL AND $4 IS NULL) OR (scope_kind = $3 AND scope_ref = $4)) + AND embedding IS NOT NULL + OFFSET 0) AS scoped + ORDER BY embedding <=> $1::vector + LIMIT $5 + """ + ) + vectorCandidateQueryEncoder + (D.rowList memoryRecordDecoder) + +recallCandidateQueryEncoder :: E.Params RecallCandidateQuery +recallCandidateQueryEncoder = + ((\q -> q.query) >$< E.param (E.nonNullable E.text)) + <> ((\q -> q.namespace) >$< E.param (E.nonNullable E.text)) + <> ((\q -> q.scopeKind) >$< E.param (E.nullable E.text)) + <> ((\q -> q.scopeRef) >$< E.param (E.nullable E.text)) + <> ((\q -> q.limit) >$< E.param (E.nonNullable E.int4)) + +vectorCandidateQueryEncoder :: E.Params VectorCandidateQuery +vectorCandidateQueryEncoder = + ((\q -> q.queryVector) >$< E.param (E.nonNullable E.text)) + <> ((\q -> q.namespace) >$< E.param (E.nonNullable E.text)) + <> ((\q -> q.scopeKind) >$< E.param (E.nullable E.text)) + <> ((\q -> q.scopeRef) >$< E.param (E.nullable E.text)) + <> ((\q -> q.limit) >$< E.param (E.nonNullable E.int4)) + +memoryRecordColumns :: Text +memoryRecordColumns = + "memory_id, agent_id, session_id, namespace, scope_kind, scope_ref, memory_type, content, priority, confidence, tags::text, status, created_at " + +memoryRecordDecoder :: D.Row MemoryRecord +memoryRecordDecoder = + makeMemoryRecord + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> (fromIntegral @Int32 @Int <$> D.column (D.nonNullable D.int4)) + <*> D.column (D.nonNullable D.text) + <*> (decodeTags <$> D.column (D.nonNullable D.text)) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.timestamptz) + +makeMemoryRecord :: + Text -> + Text -> + Maybe Text -> + Text -> + Maybe Text -> + Maybe Text -> + Text -> + Text -> + Int -> + Text -> + Set.Set Text -> + Text -> + UTCTime -> + MemoryRecord +makeMemoryRecord memoryId agentId sessionId namespace scopeKind scopeRef memoryType content priority confidence tags status createdAt = + MemoryRecord + { memoryId, + agentId, + sessionId, + scope = scopeFromColumns namespace scopeKind scopeRef, + memoryType, + content, + priority, + confidence, + tags, + status, + createdAt + } + +decodeTags :: Text -> Set.Set Text +decodeTags = + fromMaybe Set.empty . Aeson.decode . BL.fromStrict . TE.encodeUtf8 + +vectorLiteral :: Vector Double -> Text +vectorLiteral values = + "[" <> Text.intercalate "," (Text.pack . show <$> Vector.toList values) <> "]" + +-- | How many candidates each channel contributes to the RRF fusion. +-- +-- == Filtered-ANN starvation, and how it is handled +-- +-- The HNSW index is /post-filtered/: it covers the embedding column alone, so it picks its +-- candidates by distance and the namespace, scope and @status@ predicates are applied afterwards, +-- to rows it has already chosen. When the memories nearest the query sit outside the caller's +-- scope — a small scope inside a large namespace, which is the normal shape of kioku data — the +-- scan spends its entire budget on rows the filter then discards and the vector channel returns +-- __nothing__. Measured: 2000 in-scope memories and 2000 nearer ones in another namespace, at +-- default settings, returned zero rows every time. +-- +-- 'selectVectorCandidatesDiagnosed' handles this by running an exact pass whenever the +-- approximate pass comes back short of this pool. Read its Haddock for the mechanism; the summary +-- is that the approximate pass is fast and can starve, the exact pass cannot starve and costs a +-- scan of the caller's scope, and the second only runs when the first came back short. +-- +-- == Two claims that used to live here and are false +-- +-- This comment previously asserted that pgvector searches with @ef = max(ef_search, LIMIT)@, so +-- that this pool fills at the default @hnsw.ef_search@ of 40. __It does not.__ On pgvector 0.8.2, +-- a corpus of 2000 in-scope rows with /nothing/ out of scope — nothing for the filter to discard +-- at all — returned 40 rows against this LIMIT of 50. The vector channel was silently +-- under-delivering by 20% in the healthy case. 'selectVectorCandidatesDiagnosed' now sets +-- @hnsw.ef_search@ to exactly this pool size, which fills it at no measurable cost. +-- +-- It also warned against raising @hnsw.ef_search@ at all, on the evidence that +-- @SET hnsw.ef_search = 200@ flipped the planner onto an HNSW scan that starved. That evidence +-- was real but the conclusion was too broad: at 200 the planner does flip, and at 50 (this pool +-- size) it does not, while the pool fills. The distinction is measured, not argued. +-- +-- == What is still not fixed +-- +-- pgvector 0.8's @hnsw.iterative_scan@ is the vendor's own remedy for starvation and it was +-- rejected on evidence, not on principle: across five freshly built indexes on a 20000-row +-- starving corpus it returned the right answer 2 times in 5 (@relaxed_order@) and 4 times in 5 +-- (@strict_order@). HNSW construction is randomized, so it is a lottery. Do not reach for it +-- again without a sample size — a single passing run says nothing. +candidatePoolSize :: Int32 +candidatePoolSize = 50 + +rrfK :: Double +rrfK = 60 + +recencyWeight :: Double +recencyWeight = 0.10 + +prioritySignalWeight :: Double +prioritySignalWeight = 0.15 + +confidenceSignalWeight :: Double +confidenceSignalWeight = 0.05 + +recencyHalfLifeDays :: Double +recencyHalfLifeDays = 30 + +secondsPerDay :: Double +secondsPerDay = 86400 + +priorityMax :: Double +priorityMax = 100 + +alwaysInjectPriority :: Int +alwaysInjectPriority = 0 + +perMemoryCharacterBudget :: Int +perMemoryCharacterBudget = 2000 + +totalCharacterBudget :: Int +totalCharacterBudget = 12000 + +ellipsis :: Text +ellipsis = "..." + +-- | Active memories carrying __exactly__ this scope. +-- +-- Recall searches namespace-wide for a global scope; scoped reads are exact-scope. So +-- 'ScopeGlobal' here means "the rows recorded with no entity scope", /not/ "everything in the +-- namespace" — a memory under @mori:repo:web@ is returned by 'recall' with scope @mori@ but +-- not by this. For the read-side equivalent of recall's breadth, use 'getActiveInNamespace'. +getActiveByScope :: + (IOE :> es, Store :> es) => + MemoryScope -> + Eff es (Either ReadModelError [MemoryRecord]) +getActiveByScope scope = + runQueryWith + Nothing + Eventual + memoriesByScopeReadModel + (MemoriesByScopeQuery (scopeNamespaceText scope) (scopeKindText scope) (scopeRefText scope)) + +-- | Every active memory in the namespace, whatever its scope. This is the read-side +-- equivalent of what 'recall' does with a global scope. +getActiveInNamespace :: + (IOE :> es, Store :> es) => + Namespace -> + Eff es (Either ReadModelError [MemoryRecord]) +getActiveInNamespace (Namespace ns) = + runQueryWith Nothing Eventual memoriesByNamespaceReadModel (MemoriesByNamespaceQuery ns) + +-- | The global bucket of a namespace: rows recorded with no entity scope. Not the same as a +-- 'recall' scoped to the namespace, which also returns entity-scoped rows. +getGlobal :: + (IOE :> es, Store :> es) => + Namespace -> + Eff es (Either ReadModelError [MemoryRecord]) +getGlobal ns = + getActiveByScope (ScopeGlobal ns) + +getById :: + (IOE :> es, Store :> es) => + MemoryId -> + Eff es (Either ReadModelError (Maybe MemoryRecord)) +getById mid = + fmap (fmap (fmap memoryRowToRecord)) $ + runQueryWith Nothing Eventual memoryByIdReadModel (MemoryByIdQuery (idText mid)) + +getBySession :: + (IOE :> es, Store :> es) => + SessionId -> + Eff es (Either ReadModelError [MemoryRecord]) +getBySession sid = + runQueryWith Nothing Eventual memoriesBySessionReadModel (MemoriesBySessionQuery (idText sid)) + +getByType :: + (IOE :> es, Store :> es) => + Namespace -> + MemoryType -> + Eff es (Either ReadModelError [MemoryRecord]) +getByType (Namespace ns) mt = + runQueryWith Nothing Eventual memoriesByTypeReadModel (MemoriesByTypeQuery ns (memoryTypeToText mt)) + +memoryRowToRecord :: MemoryRow -> MemoryRecord +memoryRowToRecord row = + MemoryRecord + { memoryId = row.memoryId, + agentId = row.agentId, + sessionId = row.sessionId, + scope = scopeFromColumns row.namespace row.scopeKind row.scopeRef, + memoryType = row.memoryType, + content = row.content, + priority = row.priority, + confidence = row.confidence, + tags = row.tags, + status = row.status, + createdAt = row.createdAt + } diff --git a/src/Kioku/Recall/Capability.hs b/src/Kioku/Recall/Capability.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Recall/Capability.hs @@ -0,0 +1,129 @@ +module Kioku.Recall.Capability + ( VectorCapability (..), + detectVectorCapability, + classifyProbe, + CapabilityProbe (..), + ) +where + +import Data.Int (Int32) +import Effectful (Eff, (:>)) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Kioku.Prelude +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Transaction (runTransaction) + +data VectorCapability + = VectorAvailable + | VectorExtensionUnavailable + | VectorColumnsUnavailable ![Text] + | -- | @KIOKU_EMBEDDING_DIMENSIONS@ disagrees with the declared width of the + -- @kioku_memories.embedding@ column: configured first, actual second. Every embedding + -- write would fail on the @::vector@ cast, one event at a time, forever. + VectorDimensionMismatch !Int !Int + deriving stock (Generic, Eq, Show) + +data CapabilityProbe = CapabilityProbe + { hasVectorType :: !Bool, + hasEmbedding :: !Bool, + hasEmbeddingModel :: !Bool, + hasDimensions :: !Bool, + hasContentHash :: !Bool, + embeddingTypmod :: !(Maybe Int32) + } + deriving stock (Generic, Eq, Show) + +-- | Probe what the vector path can actually do, given the dimension count the process is +-- configured with (@EmbeddingConfig.dimensions@). +detectVectorCapability :: + (Store :> es) => + Int -> + Eff es VectorCapability +detectVectorCapability configuredDimensions = + classifyProbe configuredDimensions <$> runTransaction (Tx.statement () detectVectorCapabilityStmt) + +classifyProbe :: Int -> CapabilityProbe -> VectorCapability +classifyProbe configuredDimensions probe + | not probe.hasVectorType = VectorExtensionUnavailable + | not (null missing) = VectorColumnsUnavailable missing + | Just actual <- declaredDimensions, + actual /= configuredDimensions = + VectorDimensionMismatch configuredDimensions actual + | otherwise = VectorAvailable + where + missing = + missingIf (not probe.hasEmbedding) "embedding" + <> missingIf (not probe.hasEmbeddingModel) "embedding_model" + <> missingIf (not probe.hasDimensions) "dimensions" + <> missingIf (not probe.hasContentHash) "content_hash" + + -- For pgvector, a column's atttypmod *is* its declared dimension count. A column + -- declared without one reports -1, which constrains nothing, so there is nothing to + -- disagree with. + declaredDimensions = + case probe.embeddingTypmod of + Just typmod | typmod > 0 -> Just (fromIntegral typmod) + _ -> Nothing + +missingIf :: Bool -> Text -> [Text] +missingIf True columnName = [columnName] +missingIf False _ = [] + +-- | The extension check asks @to_regtype@, not @pg_extension@, and the difference matters. +-- @pg_extension@ answers "is pgvector installed /somewhere/ in this database" — but recall +-- casts with a bare @$1::vector@, and the store connects with +-- @search_path = kiroku, pg_catalog@. An extension installed into @public@ (the usual +-- operator default) satisfies @pg_extension@ and still cannot be named, so every vector +-- query would fail with @42704@ while capability detection reported everything healthy. +-- @to_regtype@ resolves against the live @search_path@, which is exactly the question the +-- query asks. +detectVectorCapabilityStmt :: Statement () CapabilityProbe +detectVectorCapabilityStmt = + preparable + """ + SELECT + to_regtype('vector') IS NOT NULL AS has_vector_type, + EXISTS ( + SELECT 1 + FROM information_schema.columns + WHERE table_schema = 'kiroku' AND table_name = 'kioku_memories' AND column_name = 'embedding' + ) AS has_embedding, + EXISTS ( + SELECT 1 + FROM information_schema.columns + WHERE table_schema = 'kiroku' AND table_name = 'kioku_memories' AND column_name = 'embedding_model' + ) AS has_embedding_model, + EXISTS ( + SELECT 1 + FROM information_schema.columns + WHERE table_schema = 'kiroku' AND table_name = 'kioku_memories' AND column_name = 'dimensions' + ) AS has_dimensions, + EXISTS ( + SELECT 1 + FROM information_schema.columns + WHERE table_schema = 'kiroku' AND table_name = 'kioku_memories' AND column_name = 'content_hash' + ) AS has_content_hash, + ( + SELECT a.atttypmod + FROM pg_attribute a + JOIN pg_class c ON c.oid = a.attrelid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = 'kiroku' + AND c.relname = 'kioku_memories' + AND a.attname = 'embedding' + AND NOT a.attisdropped + ) AS embedding_typmod + """ + E.noParams + ( D.singleRow $ + CapabilityProbe + <$> D.column (D.nonNullable D.bool) + <*> D.column (D.nonNullable D.bool) + <*> D.column (D.nonNullable D.bool) + <*> D.column (D.nonNullable D.bool) + <*> D.column (D.nonNullable D.bool) + <*> D.column (D.nullable D.int4) + ) diff --git a/src/Kioku/Session.hs b/src/Kioku/Session.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Session.hs @@ -0,0 +1,542 @@ +module Kioku.Session + ( SessionRow (..), + SessionWriteError (..), + start, + awaitInput, + resume, + forceResume, + complete, + failSession, + recordInteractive, + recordTurn, + getById, + getRecentInNamespace, + getByScope, + getByFocus, + getByStartedRange, + getChain, + getDelegationChildren, + getAwaitingByCorrelationKey, + getTurns, + ) +where + +import Data.List (find) +import Data.Text qualified as Text +import Effectful (Eff, IOE, (:>)) +import Effectful.Error.Static (Error) +import Keiro.Command (CommandError (..), defaultRunCommandOptions) +import Keiro.Projection (runCommandWithProjections) +import Keiro.ReadModel (ConsistencyMode (..), ReadModelError, runQueryWith) +import Kioku.Api.Scope (MemoryScope, Namespace (..), scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Distill.Timer (l1TimerScheduleProjection) +import Kioku.Id (SessionId, idText) +import Kioku.Prelude +import Kioku.Session.Domain +import Kioku.Session.EventStream (sessionEventStream, sessionStream) +import Kioku.Session.ReadModel + ( AwaitingSessionsByCorrelationKeyQuery (..), + SessionByIdQuery (..), + SessionChainQuery (..), + SessionDelegationChildrenQuery (..), + SessionRow (..), + SessionsByFocusQuery (..), + SessionsByNamespaceQuery (..), + SessionsByScopeQuery (..), + SessionsByStartedRangeQuery (..), + TurnRow (..), + TurnsBySessionQuery (..), + awaitingSessionsByCorrelationKeyReadModel, + sessionByIdReadModel, + sessionChainReadModel, + sessionDelegationChildrenReadModel, + sessionInlineProjection, + sessionsByFocusReadModel, + sessionsByNamespaceReadModel, + sessionsByScopeReadModel, + sessionsByStartedRangeReadModel, + turnsBySessionReadModel, + ) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) + +data SessionWriteError + = SessionCommandRejected !CommandError + | SessionReadFailed !ReadModelError + | SessionNotFound + | SessionNotRunning + | SessionNotAwaiting + | SessionCorrelationMismatch + | SessionInvalidLineage !Text + | SessionConflict !Text + deriving stock (Generic, Show) + +-- | The deepest delegation chain a session may declare. Far above any legitimate agent +-- hierarchy; it exists to bound absurd input, not to express a product limit. +maxDelegationDepth :: Int +maxDelegationDepth = 64 + +start :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + StartSessionData -> + Eff es (Either SessionWriteError SessionId) +start cmdData = + case validateLineage cmdData of + Just reason -> pure (Left (SessionInvalidLineage reason)) + Nothing -> do + existing <- getById cmdData.sessionId + case existing of + Left err -> pure (Left (SessionReadFailed err)) + Right (Just row) -> pure (idempotentOr "start" startMismatch row cmdData.sessionId) + Right Nothing -> + runSessionCommand cmdData.sessionId (StartSession cmdData) + >>= acceptRejectedIfMatches cmdData.sessionId (isNothing . startMismatch) + where + startMismatch = mismatchOf sessionStartFields cmdData + +-- | Pure, command-time lineage checks: 'Just' a reason means reject. +-- +-- Deliberately does /not/ check that the referenced sessions exist. Existence checks would +-- be read-model reads with their own races, and they would forbid legitimate out-of-order +-- ingestion; a dangling pointer is harmless to the chain query, which simply stops walking. +-- Because existence is unchecked a cycle can still be constructed (write A→B before B +-- exists, then B→A), which is why 'selectSessionChainStmt' carries its own cycle guard. +-- +-- These run only here, never during replay, so no historical event can be made +-- unreplayable by tightening them. +validateLineage :: StartSessionData -> Maybe Text +validateLineage d + | d.previousSessionId == Just d.sessionId = + Just "previousSessionId must not be the session's own id" + | d.parentSessionId == Just d.sessionId = + Just "parentSessionId must not be the session's own id" + | d.delegationDepth < 0 = + Just "delegationDepth must not be negative" + | d.delegationDepth > maxDelegationDepth = + Just ("delegationDepth must not exceed " <> Text.pack (show maxDelegationDepth)) + | isJust d.parentSessionId && d.delegationDepth < 1 = + Just "a delegated session (parentSessionId present) must have delegationDepth >= 1" + | isNothing d.parentSessionId && d.delegationDepth /= 0 = + Just "a root session (no parentSessionId) must have delegationDepth 0" + | otherwise = Nothing + +-- * Idempotent accepts + +-- | The five statuses the projection writes. Parsing at the point of decision keeps +-- 'SessionRow.status' a 'Text' (it is part of the read-model shape hosts consume) while +-- removing stringly-typed comparisons from the decision logic. +data SessionStatus + = StatusRunning + | StatusAwaiting + | StatusCompleted + | StatusFailed + | StatusInteractive + deriving stock (Eq, Show) + +parseSessionStatus :: Text -> Maybe SessionStatus +parseSessionStatus = \case + "running" -> Just StatusRunning + "awaiting" -> Just StatusAwaiting + "completed" -> Just StatusCompleted + "failed" -> Just StatusFailed + "interactive" -> Just StatusInteractive + _ -> Nothing + +-- | Look up the session and hand its row plus parsed status to the caller. +withExistingSession :: + (IOE :> es, Store :> es) => + SessionId -> + (SessionRow -> SessionStatus -> Eff es (Either SessionWriteError SessionId)) -> + Eff es (Either SessionWriteError SessionId) +withExistingSession sid k = do + existing <- getById sid + case existing of + Left err -> pure (Left (SessionReadFailed err)) + Right Nothing -> pure (Left SessionNotFound) + Right (Just row) -> + case parseSessionStatus row.status of + Nothing -> pure (Left (SessionConflict ("unrecognized session status: " <> row.status))) + Just status -> k row status + +-- | A named comparison between one request field and the row that already exists. +type FieldCheck cmd = (Text, cmd -> SessionRow -> Bool) + +-- | The first request field that disagrees with the recorded row, if any. +-- +-- Call-time timestamps (@startedAt@, @completedAt@, @failedAt@, @resumedAt@) are +-- deliberately /not/ compared. The session id is the identity, so a second write against it +-- carrying the same semantic payload is a retry — and a retry that re-reads the clock is +-- the normal shape of one. Comparing the timestamp would turn every such retry into a hard +-- conflict; kioku's own distillation pass does exactly this on the memory side (see +-- 'Kioku.Memory.mismatchOf'), which is what proved the point. +-- +-- Semantic payload is compared, including 'awaitingDeadline' — a deadline is something the +-- caller /asked for/, not a record of when it called. +mismatchOf :: [FieldCheck cmd] -> cmd -> SessionRow -> Maybe Text +mismatchOf checks cmd row = + fst <$> find (\(_, matches) -> not (matches cmd row)) checks + +-- | A duplicate request that matches what already happened succeeds; one that conflicts +-- with it gets a conflict error naming the field that differs. +idempotentOr :: + Text -> + (SessionRow -> Maybe Text) -> + SessionRow -> + SessionId -> + Either SessionWriteError SessionId +idempotentOr operation mismatch row sid = + case mismatch row of + Nothing -> Right sid + Just field -> + Left (SessionConflict (operation <> ": " <> field <> " differs from the recorded session")) + +-- | Translate a losing race into the success the winner got. +-- +-- Two identical requests can be in flight at once; keiro's optimistic-concurrency retry +-- lets one win and rejects the other. Re-reading the row after a rejection tells us which +-- kind of loser this is: if the observed state now matches what we asked for, the write we +-- wanted happened (someone else did it) and the caller gets the idempotent success. A +-- genuinely conflicting loser still gets its rejection. +acceptRejectedIfMatches :: + (IOE :> es, Store :> es) => + SessionId -> + (SessionRow -> Bool) -> + Either SessionWriteError SessionId -> + Eff es (Either SessionWriteError SessionId) +acceptRejectedIfMatches sid matches = \case + Left err@(SessionCommandRejected CommandRejected) -> do + reread <- getById sid + pure case reread of + Right (Just row) | matches row -> Right sid + _ -> Left err + other -> pure other + +sessionStartFields :: [FieldCheck StartSessionData] +sessionStartFields = + [ ("agentId", \d row -> row.agentId == d.agentId), + ("focus", \d row -> row.focus == d.focus), + ("namespace", \d row -> row.namespace == scopeNamespaceText d.scope), + ("scopeKind", \d row -> row.scopeKind == scopeKindText d.scope), + ("scopeRef", \d row -> row.scopeRef == scopeRefText d.scope), + ("subjectRef", \d row -> row.subjectRef == d.subjectRef), + ("previousSessionId", \d row -> row.previousSessionId == (idText <$> d.previousSessionId)), + ("parentSessionId", \d row -> row.parentSessionId == (idText <$> d.parentSessionId)), + ("delegationDepth", \d row -> row.delegationDepth == d.delegationDepth) + ] + +sessionInteractiveFields :: [FieldCheck RecordInteractiveSessionData] +sessionInteractiveFields = + [ ("status", \_ row -> parseSessionStatus row.status == Just StatusInteractive), + ("agentId", \d row -> row.agentId == d.agentId), + ("focus", \d row -> row.focus == d.focus), + ("namespace", \d row -> row.namespace == scopeNamespaceText d.scope), + ("scopeKind", \d row -> row.scopeKind == scopeKindText d.scope), + ("scopeRef", \d row -> row.scopeRef == scopeRefText d.scope), + ("subjectRef", \d row -> row.subjectRef == d.subjectRef) + ] + +sessionAwaitFields :: [FieldCheck AwaitInputData] +sessionAwaitFields = + [ ("awaitingReason", \d row -> row.awaitingReason == Just d.reason), + ("awaitingCorrelationKey", \d row -> row.awaitingCorrelationKey == d.correlationKey), + ("awaitingDeadline", \d row -> row.awaitingDeadline == d.deadline) + ] + +sessionResumeFields :: [FieldCheck ResumeSessionData] +sessionResumeFields = + [("resumeInput", \d row -> row.resumeInput == Just d.input)] + +sessionCompleteFields :: [FieldCheck CompleteSessionData] +sessionCompleteFields = + [ ("modelUsed", \d row -> row.modelUsed == d.modelUsed), + ("summary", \d row -> row.summary == d.summary) + ] + +sessionFailFields :: [FieldCheck FailSessionData] +sessionFailFields = + [("errorMessage", \d row -> row.errorMessage == Just d.errorMessage)] + +complete :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + CompleteSessionData -> + Eff es (Either SessionWriteError SessionId) +complete cmdData = + withExistingSession cmdData.sessionId \row status -> + case status of + StatusRunning -> runComplete + StatusAwaiting -> runComplete + StatusCompleted -> pure (idempotentOr "complete" completeMismatch row cmdData.sessionId) + StatusFailed -> pure (Left (SessionConflict "complete: the session already failed")) + StatusInteractive -> + pure (Left (SessionConflict "complete: an interactive session has no lifecycle to complete")) + where + completeMismatch = mismatchOf sessionCompleteFields cmdData + runComplete = + runSessionCommand cmdData.sessionId (CompleteSession cmdData) + >>= acceptRejectedIfMatches cmdData.sessionId (isNothing . completeMismatch) + +failSession :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + FailSessionData -> + Eff es (Either SessionWriteError SessionId) +failSession cmdData = + withExistingSession cmdData.sessionId \row status -> + case status of + StatusRunning -> runFail + StatusAwaiting -> runFail + StatusFailed -> pure (idempotentOr "failSession" failMismatch row cmdData.sessionId) + StatusCompleted -> + pure (Left (SessionConflict "failSession: the session already completed successfully")) + StatusInteractive -> + pure (Left (SessionConflict "failSession: an interactive session has no lifecycle to fail")) + where + failMismatch = mismatchOf sessionFailFields cmdData + runFail = + runSessionCommand cmdData.sessionId (FailSession cmdData) + >>= acceptRejectedIfMatches cmdData.sessionId (isNothing . failMismatch) + +awaitInput :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + AwaitInputData -> + Eff es (Either SessionWriteError SessionId) +awaitInput cmdData = + withExistingSession cmdData.sessionId \row status -> + case status of + StatusAwaiting -> pure (idempotentOr "awaitInput" awaitMismatch row cmdData.sessionId) + StatusRunning -> + runSessionCommand cmdData.sessionId (AwaitInput cmdData) + >>= acceptRejectedIfMatches cmdData.sessionId (isNothing . awaitMismatch) + _ -> pure (Left SessionNotRunning) + where + awaitMismatch = mismatchOf sessionAwaitFields cmdData + +-- | Resume a parked session with the key it parked on. +-- +-- The correlation key must match exactly — a keyed resume of a keyless wait, or of a wait +-- on a different key, is rejected. An omitted key no longer bypasses matching; use +-- 'forceResume' for that, explicitly. +-- +-- The precheck below only shapes a friendly early error. The real enforcement is the +-- aggregate's own guard, which keiro re-evaluates after any optimistic-concurrency retry — +-- so a stale caller cannot resume a wait that was already resumed and re-parked. +resume :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + ResumeSessionData -> + Eff es (Either SessionWriteError SessionId) +resume cmdData = + withExistingSession cmdData.sessionId \row status -> + case status of + -- Already running: a re-delivery of *this* resume is a success; a different input + -- means someone else answered the wait, which is a conflict, not an idempotent hit. + StatusRunning -> pure (idempotentOr "resume" resumeMismatch row cmdData.sessionId) + StatusAwaiting + | not cmdData.force && row.awaitingCorrelationKey /= cmdData.correlationKey -> + pure (Left SessionCorrelationMismatch) + | otherwise -> + runSessionCommand cmdData.sessionId (ResumeSession cmdData) + >>= acceptRejectedIfMatches cmdData.sessionId (isNothing . resumeMismatch) + _ -> pure (Left SessionNotAwaiting) + where + resumeMismatch = mismatchOf sessionResumeFields cmdData + +-- | Resume a parked session regardless of which key it parked on. +-- +-- An operator/host override for unsticking a session whose awaited key is lost or wrong. +-- It is inherently last-writer-wins: if the session is concurrently re-parked on a new +-- wait, a force resume may answer the wrong one. Prefer 'resume'. +forceResume :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + Text -> + UTCTime -> + Eff es (Either SessionWriteError SessionId) +forceResume sid input resumedAt = + resume + ResumeSessionData + { sessionId = sid, + correlationKey = Nothing, + force = True, + input, + resumedAt + } + +recordInteractive :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + RecordInteractiveSessionData -> + Eff es (Either SessionWriteError SessionId) +recordInteractive cmdData = do + existing <- getById cmdData.sessionId + case existing of + Left err -> pure (Left (SessionReadFailed err)) + Right (Just row) -> pure (idempotentOr "recordInteractive" interactiveMismatch row cmdData.sessionId) + Right Nothing -> + runSessionCommand cmdData.sessionId (RecordInteractiveSession cmdData) + >>= acceptRejectedIfMatches cmdData.sessionId (isNothing . interactiveMismatch) + where + interactiveMismatch = mismatchOf sessionInteractiveFields cmdData + +-- | Record one turn of a running session. +-- +-- Turn identity is @(sessionId, turnIndex)@; @turnId@ is an idempotency token that travels +-- with it. A re-delivery of an identical turn succeeds without appending an event; the same +-- index carrying different content, or the same @turnId@ reappearing at a different index, +-- is a conflict. The aggregate independently enforces that indexes strictly increase. +-- +-- Reusing a @turnId@ across two /different sessions/ remains a raw primary-key violation +-- surfaced as @StoreFailed@: turn ids are host-generated, so a cross-session collision is a +-- caller bug, and mapping that specific SQL error from inside keiro's projection +-- transaction is not worth the machinery. +recordTurn :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + RecordTurnData -> + Eff es (Either SessionWriteError SessionId) +recordTurn cmdData = + withExistingSession cmdData.sessionId \_row status -> + case status of + StatusRunning -> do + turns <- getTurns cmdData.sessionId + case turns of + Left err -> pure (Left (SessionReadFailed err)) + Right existingTurns -> + case turnVerdict cmdData existingTurns of + Just verdict -> pure verdict + Nothing -> + runSessionCommand cmdData.sessionId (RecordTurn cmdData) + >>= acceptRejectedTurnIfMatches cmdData + _ -> pure (Left SessionNotRunning) + +-- | The turns-table counterpart of 'acceptRejectedIfMatches': a concurrent duplicate that +-- lost the optimistic-concurrency race converges to the winner's success. +acceptRejectedTurnIfMatches :: + (IOE :> es, Store :> es) => + RecordTurnData -> + Either SessionWriteError SessionId -> + Eff es (Either SessionWriteError SessionId) +acceptRejectedTurnIfMatches d = \case + Left err@(SessionCommandRejected CommandRejected) -> do + turns <- getTurns d.sessionId + pure case turns of + Right rows + | Just row <- find (\row -> row.turnIndex == d.turnIndex) rows, + turnRowMatches d row -> + Right d.sessionId + _ -> Left err + other -> pure other + +-- | 'Just' a final answer if an existing turn row already decides this request; +-- 'Nothing' if the command should run. +turnVerdict :: RecordTurnData -> [TurnRow] -> Maybe (Either SessionWriteError SessionId) +turnVerdict d existingTurns + | Just row <- sameIndex = + Just + if turnRowMatches d row + then Right d.sessionId + else Left (SessionConflict ("recordTurn: turn " <> Text.pack (show d.turnIndex) <> " already recorded with different content")) + | any (\row -> row.turnId == d.turnId) existingTurns = + Just (Left (SessionConflict ("recordTurn: turnId " <> d.turnId <> " is already used at a different turn index"))) + | otherwise = Nothing + where + sameIndex = find (\row -> row.turnIndex == d.turnIndex) existingTurns + +turnRowMatches :: RecordTurnData -> TurnRow -> Bool +turnRowMatches d row = + row.turnId == d.turnId + && row.role == d.role + && row.content == d.content + && row.toolSummary == d.toolSummary + && row.promptTokens == d.promptTokens + && row.outputTokens == d.outputTokens + +getById :: + (IOE :> es, Store :> es) => + SessionId -> + Eff es (Either ReadModelError (Maybe SessionRow)) +getById sid = + runQueryWith Nothing Eventual sessionByIdReadModel (SessionByIdQuery (idText sid)) + +getRecentInNamespace :: + (IOE :> es, Store :> es) => + Namespace -> + Int -> + Eff es (Either ReadModelError [SessionRow]) +getRecentInNamespace ns limit = + runQueryWith Nothing Eventual sessionsByNamespaceReadModel (SessionsByNamespaceQuery (namespaceText ns) limit) + +getByScope :: + (IOE :> es, Store :> es) => + MemoryScope -> + Eff es (Either ReadModelError [SessionRow]) +getByScope scope = + runQueryWith + Nothing + Eventual + sessionsByScopeReadModel + (SessionsByScopeQuery (scopeNamespaceText scope) (scopeKindText scope) (scopeRefText scope)) + +getByFocus :: + (IOE :> es, Store :> es) => + Namespace -> + Text -> + Eff es (Either ReadModelError [SessionRow]) +getByFocus ns focus = + runQueryWith Nothing Eventual sessionsByFocusReadModel (SessionsByFocusQuery (namespaceText ns) focus) + +getByStartedRange :: + (IOE :> es, Store :> es) => + Namespace -> + UTCTime -> + UTCTime -> + Eff es (Either ReadModelError [SessionRow]) +getByStartedRange ns startedAfter startedBefore = + runQueryWith Nothing Eventual sessionsByStartedRangeReadModel (SessionsByStartedRangeQuery (namespaceText ns) startedAfter startedBefore) + +getChain :: + (IOE :> es, Store :> es) => + SessionId -> + Eff es (Either ReadModelError [SessionRow]) +getChain sid = + runQueryWith Nothing Eventual sessionChainReadModel (SessionChainQuery (idText sid)) + +getDelegationChildren :: + (IOE :> es, Store :> es) => + SessionId -> + Eff es (Either ReadModelError [SessionRow]) +getDelegationChildren sid = + runQueryWith Nothing Eventual sessionDelegationChildrenReadModel (SessionDelegationChildrenQuery (idText sid)) + +getAwaitingByCorrelationKey :: + (IOE :> es, Store :> es) => + Namespace -> + Text -> + Eff es (Either ReadModelError [SessionRow]) +getAwaitingByCorrelationKey ns correlationKey = + runQueryWith Nothing Eventual awaitingSessionsByCorrelationKeyReadModel (AwaitingSessionsByCorrelationKeyQuery (namespaceText ns) correlationKey) + +namespaceText :: Namespace -> Text +namespaceText (Namespace ns) = ns + +getTurns :: + (IOE :> es, Store :> es) => + SessionId -> + Eff es (Either ReadModelError [TurnRow]) +getTurns sid = + runQueryWith Nothing Eventual turnsBySessionReadModel (TurnsBySessionQuery (idText sid)) + +runSessionCommand :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + SessionCommand -> + Eff es (Either SessionWriteError SessionId) +runSessionCommand sid cmd = do + result <- + runCommandWithProjections + defaultRunCommandOptions + sessionEventStream + (sessionStream sid) + cmd + [sessionInlineProjection, l1TimerScheduleProjection] + pure $ + case result of + Left err -> Left (SessionCommandRejected err) + Right _ -> Right sid diff --git a/src/Kioku/Session/Domain.hs b/src/Kioku/Session/Domain.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Session/Domain.hs @@ -0,0 +1,444 @@ +{-# LANGUAGE TemplateHaskell #-} + +module Kioku.Session.Domain + ( SessionVertex (..), + SessionRegs, + StartSessionData (..), + CompleteSessionData (..), + FailSessionData (..), + Continuation (..), + AwaitInputData (..), + ResumeSessionData (..), + RecordInteractiveSessionData (..), + RecordTurnData (..), + SessionCommand (..), + commandSessionId, + SessionStartedData (..), + SessionCompletedData (..), + SessionFailedData (..), + SessionAwaitingData (..), + SessionResumedData (..), + InteractiveSessionRecordedData (..), + TurnRecordedData (..), + SessionEvent (..), + eventSessionId, + sessionTransducer, + ) +where + +import Data.Aeson.Types (withObject, (.!=), (.:), (.:?)) +import Keiki.Builder ((=:)) +import Keiki.Builder qualified as B +import Keiki.Core (HsPred, SymTransducer, lit, (.==), (.||)) +import Keiki.Generics (emptyRegFile) +import Keiki.Generics.TH (deriveAggregate) +import Kioku.Api.Scope (MemoryScope) +import Kioku.Id (SessionId) +import Kioku.Prelude + +data SessionVertex = NotCreated | Running | Completed | Failed | Interactive | Awaiting + deriving stock (Eq, Ord, Show, Enum, Bounded) + +-- | Replayed aggregate state carried alongside the vertex. +-- +-- @awaitedCorrelationKey@ is the key the session is currently parked on (set by +-- @SessionAwaiting@, cleared by @SessionResumed@). It is what makes resume-correlation +-- matching an aggregate invariant rather than a racy read-model precheck. +-- +-- @lastTurnIndex@ is the highest turn index committed so far (-1 before any turn), which +-- makes @RecordTurn@'s strictly-increasing index contract enforceable in the state machine +-- rather than only at the command layer. +type SessionRegs = + '[ '("awaitedCorrelationKey", Maybe Text), + '("lastTurnIndex", Int) + ] + +data StartSessionData = StartSessionData + { sessionId :: !SessionId, + agentId :: !Text, + focus :: !Text, + scope :: !MemoryScope, + subjectRef :: !(Maybe Text), + previousSessionId :: !(Maybe SessionId), + parentSessionId :: !(Maybe SessionId), + delegationDepth :: !Int, + startedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data CompleteSessionData = CompleteSessionData + { sessionId :: !SessionId, + completedAt :: !UTCTime, + modelUsed :: !(Maybe Text), + summary :: !(Maybe Text) + } + deriving stock (Generic, Eq, Show) + +data FailSessionData = FailSessionData + { sessionId :: !SessionId, + failedAt :: !UTCTime, + errorMessage :: !Text + } + deriving stock (Generic, Eq, Show) + +-- | What a parked session is waiting for. +-- +-- @deadline@ is advisory only: it is stored for hosts and kioku does not enforce it. No timer +-- fires and nothing expires when it passes (MasterPlan 2 decision, 2026-07-07). +data Continuation = Continuation + { reason :: !Text, + correlationKey :: !(Maybe Text), + deadline :: !(Maybe UTCTime) + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Park a session until a host supplies input. +-- +-- @deadline@ is advisory only: it is stored for hosts and kioku does not enforce it. No timer +-- fires and nothing expires when it passes (MasterPlan 2 decision, 2026-07-07). +data AwaitInputData = AwaitInputData + { sessionId :: !SessionId, + reason :: !Text, + correlationKey :: !(Maybe Text), + deadline :: !(Maybe UTCTime), + awaitedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +-- | Resume a parked session. +-- +-- @correlationKey@ must equal the key the session parked on, or the aggregate rejects the +-- command. @force@ waives that check; it is set only by 'Kioku.Session.forceResume' and is +-- inherently last-writer-wins. +data ResumeSessionData = ResumeSessionData + { sessionId :: !SessionId, + correlationKey :: !(Maybe Text), + force :: !Bool, + input :: !Text, + resumedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data RecordInteractiveSessionData = RecordInteractiveSessionData + { sessionId :: !SessionId, + agentId :: !Text, + focus :: !Text, + scope :: !MemoryScope, + subjectRef :: !(Maybe Text), + startedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data RecordTurnData = RecordTurnData + { sessionId :: !SessionId, + turnId :: !Text, + turnIndex :: !Int, + role :: !Text, + content :: !Text, + toolSummary :: !(Maybe Text), + promptTokens :: !(Maybe Int), + outputTokens :: !(Maybe Int), + recordedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data SessionCommand + = StartSession !StartSessionData + | CompleteSession !CompleteSessionData + | FailSession !FailSessionData + | AwaitInput !AwaitInputData + | ResumeSession !ResumeSessionData + | RecordInteractiveSession !RecordInteractiveSessionData + | RecordTurn !RecordTurnData + deriving stock (Generic, Eq, Show) + +commandSessionId :: SessionCommand -> SessionId +commandSessionId = \case + StartSession d -> d.sessionId + CompleteSession d -> d.sessionId + FailSession d -> d.sessionId + AwaitInput d -> d.sessionId + ResumeSession d -> d.sessionId + RecordInteractiveSession d -> d.sessionId + RecordTurn d -> d.sessionId + +data SessionStartedData = SessionStartedData + { sessionId :: !SessionId, + agentId :: !Text, + focus :: !Text, + scope :: !MemoryScope, + subjectRef :: !(Maybe Text), + previousSessionId :: !(Maybe SessionId), + parentSessionId :: !(Maybe SessionId), + delegationDepth :: !Int, + startedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToJSON) + +instance FromJSON SessionStartedData where + parseJSON = + withObject "SessionStartedData" \o -> + SessionStartedData + <$> o .: "sessionId" + <*> o .: "agentId" + <*> o .: "focus" + <*> o .: "scope" + <*> o .:? "subjectRef" + <*> o .:? "previousSessionId" + <*> o .:? "parentSessionId" + <*> o .:? "delegationDepth" .!= 0 + <*> o .: "startedAt" + +data SessionCompletedData = SessionCompletedData + { sessionId :: !SessionId, + completedAt :: !UTCTime, + modelUsed :: !(Maybe Text), + summary :: !(Maybe Text) + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data SessionFailedData = SessionFailedData + { sessionId :: !SessionId, + failedAt :: !UTCTime, + errorMessage :: !Text + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data SessionAwaitingData = SessionAwaitingData + { sessionId :: !SessionId, + reason :: !Text, + correlationKey :: !(Maybe Text), + deadline :: !(Maybe UTCTime), + awaitedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data SessionResumedData = SessionResumedData + { sessionId :: !SessionId, + correlationKey :: !(Maybe Text), + force :: !Bool, + input :: !Text, + resumedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (ToJSON) + +-- | Events written before the resume-correlation guard landed carry no @force@ key. Under +-- the old code an omitted correlation key bypassed matching entirely, so a keyless legacy +-- resume decodes as a force-resume and a keyed one as a plain resume. That is what keeps +-- every historical stream replayable through the guard the transducer now applies. +instance FromJSON SessionResumedData where + parseJSON = + withObject "SessionResumedData" \o -> do + sessionId <- o .: "sessionId" + correlationKey <- o .:? "correlationKey" + force <- o .:? "force" .!= isNothing correlationKey + input <- o .: "input" + resumedAt <- o .: "resumedAt" + pure SessionResumedData {sessionId, correlationKey, force, input, resumedAt} + +data InteractiveSessionRecordedData = InteractiveSessionRecordedData + { sessionId :: !SessionId, + agentId :: !Text, + focus :: !Text, + scope :: !MemoryScope, + subjectRef :: !(Maybe Text), + startedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data TurnRecordedData = TurnRecordedData + { sessionId :: !SessionId, + turnId :: !Text, + turnIndex :: !Int, + role :: !Text, + content :: !Text, + toolSummary :: !(Maybe Text), + promptTokens :: !(Maybe Int), + outputTokens :: !(Maybe Int), + recordedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + deriving anyclass (FromJSON, ToJSON) + +data SessionEvent + = SessionStarted !SessionStartedData + | SessionCompleted !SessionCompletedData + | SessionFailed !SessionFailedData + | SessionAwaiting !SessionAwaitingData + | SessionResumed !SessionResumedData + | InteractiveSessionRecorded !InteractiveSessionRecordedData + | TurnRecorded !TurnRecordedData + deriving stock (Generic, Eq, Show) + +instance FromJSON SessionEvent where + parseJSON = genericParseJSON eventAesonOptions + +instance ToJSON SessionEvent where + toJSON = genericToJSON eventAesonOptions + +eventSessionId :: SessionEvent -> SessionId +eventSessionId = \case + SessionStarted d -> d.sessionId + SessionCompleted d -> d.sessionId + SessionFailed d -> d.sessionId + SessionAwaiting d -> d.sessionId + SessionResumed d -> d.sessionId + InteractiveSessionRecorded d -> d.sessionId + TurnRecorded d -> d.sessionId + +$(deriveAggregate ''SessionCommand ''SessionRegs ''SessionEvent) + +sessionTransducer :: + SymTransducer + (HsPred SessionRegs SessionCommand) + SessionRegs + SessionVertex + SessionCommand + SessionEvent +sessionTransducer = + B.buildTransducer NotCreated emptyRegFile isTerminal do + B.from NotCreated do + B.onCmd inCtorStartSession $ \d -> B.do + -- 'emptyRegFile' binds every slot to a deferred error, so this edge — the only way + -- into Running, and thus into Awaiting and RecordTurn — must initialize both. + B.slot @"awaitedCorrelationKey" =: lit Nothing + B.slot @"lastTurnIndex" =: lit (-1) + B.emit + wireSessionStarted + SessionStartedTermFields + { sessionId = d.sessionId, + agentId = d.agentId, + focus = d.focus, + scope = d.scope, + subjectRef = d.subjectRef, + previousSessionId = d.previousSessionId, + parentSessionId = d.parentSessionId, + delegationDepth = d.delegationDepth, + startedAt = d.startedAt + } + B.goto Running + + B.onCmd inCtorRecordInteractiveSession $ \d -> B.do + B.emit + wireInteractiveSessionRecorded + InteractiveSessionRecordedTermFields + { sessionId = d.sessionId, + agentId = d.agentId, + focus = d.focus, + scope = d.scope, + subjectRef = d.subjectRef, + startedAt = d.startedAt + } + B.goto Interactive + + B.from Running do + B.onCmd inCtorCompleteSession $ \d -> B.do + B.emit + wireSessionCompleted + SessionCompletedTermFields + { sessionId = d.sessionId, + completedAt = d.completedAt, + modelUsed = d.modelUsed, + summary = d.summary + } + B.goto Completed + + B.onCmd inCtorFailSession $ \d -> B.do + B.emit + wireSessionFailed + SessionFailedTermFields + { sessionId = d.sessionId, + failedAt = d.failedAt, + errorMessage = d.errorMessage + } + B.goto Failed + + B.onCmd inCtorRecordTurn $ \d -> B.do + -- Turn identity: (sessionId, turnIndex). Indexes must strictly increase, so a + -- re-delivered or out-of-order turn cannot silently overwrite a committed one. + -- 'turnIndex' is already in the event payload, so replay recovers it and existing + -- strictly-increasing streams rehydrate unchanged (verified by Audit B). + B.requireGt d.turnIndex (B.reg @"lastTurnIndex") + B.slot @"lastTurnIndex" =: d.turnIndex + B.emit + wireTurnRecorded + TurnRecordedTermFields + { sessionId = d.sessionId, + turnId = d.turnId, + turnIndex = d.turnIndex, + role = d.role, + content = d.content, + toolSummary = d.toolSummary, + promptTokens = d.promptTokens, + outputTokens = d.outputTokens, + recordedAt = d.recordedAt + } + B.goto Running + + B.onCmd inCtorAwaitInput $ \d -> B.do + B.slot @"awaitedCorrelationKey" =: d.correlationKey + B.emit + wireSessionAwaiting + SessionAwaitingTermFields + { sessionId = d.sessionId, + reason = d.reason, + correlationKey = d.correlationKey, + deadline = d.deadline, + awaitedAt = d.awaitedAt + } + B.goto Awaiting + + B.from Awaiting do + B.onCmd inCtorResumeSession $ \d -> B.do + -- The resume must name the key this session actually parked on, unless it is an + -- explicit force. Enforcing it here rather than in a read-model precheck is what + -- closes the race: keiro re-runs this edge against the post-conflict state. + B.requireGuard + ((d.force .== lit True) .|| (d.correlationKey .== B.reg @"awaitedCorrelationKey")) + B.slot @"awaitedCorrelationKey" =: lit Nothing + B.emit + wireSessionResumed + SessionResumedTermFields + { sessionId = d.sessionId, + correlationKey = d.correlationKey, + -- Mandatory for replay: the guard reads 'force', and hydration can only + -- recover command fields that the event payload carries. + force = d.force, + input = d.input, + resumedAt = d.resumedAt + } + B.goto Running + + B.onCmd inCtorCompleteSession $ \d -> B.do + B.emit + wireSessionCompleted + SessionCompletedTermFields + { sessionId = d.sessionId, + completedAt = d.completedAt, + modelUsed = d.modelUsed, + summary = d.summary + } + B.goto Completed + + B.onCmd inCtorFailSession $ \d -> B.do + B.emit + wireSessionFailed + SessionFailedTermFields + { sessionId = d.sessionId, + failedAt = d.failedAt, + errorMessage = d.errorMessage + } + B.goto Failed + where + isTerminal = \case + Completed -> True + Failed -> True + Interactive -> True + _ -> False diff --git a/src/Kioku/Session/EventStream.hs b/src/Kioku/Session/EventStream.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Session/EventStream.hs @@ -0,0 +1,169 @@ +module Kioku.Session.EventStream + ( SessionEventStream, + sessionEventStream, + sessionCodec, + sessionStream, + parseSessionEvent, + ) +where + +import Data.Aeson (Value) +import Data.Aeson.Types (Parser, parseEither, withObject, (.:), (.:?)) +import Data.Text qualified as Text +import Keiki.Core (HsPred) +import Keiki.Generics (emptyRegFile) +import Keiro.Codec (Codec (..), EventType (..)) +import Keiro.EventStream (EventStream (..), SnapshotPolicy (..)) +import Keiro.EventStream.Validate (ValidatedEventStream, mkEventStreamOrThrow) +import Keiro.Stream (Stream) +import Keiro.Stream qualified as Stream +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), ScopeKind (..)) +import Kioku.Id (SessionId, idText, parseIdLenient) +import Kioku.Prelude +import Kioku.Session.Domain + +type SessionEventStream = + EventStream (HsPred SessionRegs SessionCommand) SessionRegs SessionVertex SessionCommand SessionEvent + +sessionStream :: SessionId -> Stream SessionEventStream +sessionStream sid = Stream.entityStream (Stream.categoryUnsafe "kioku_session") (idText sid) + +sessionEventStream :: ValidatedEventStream (HsPred SessionRegs SessionCommand) SessionRegs SessionVertex SessionCommand SessionEvent +sessionEventStream = + mkEventStreamOrThrow "kioku-session" sessionEventStreamDefinition + +sessionEventStreamDefinition :: SessionEventStream +sessionEventStreamDefinition = + EventStream + { transducer = sessionTransducer, + initialState = NotCreated, + initialRegisters = emptyRegFile, + eventCodec = sessionCodec, + resolveStreamName = Stream.streamName, + snapshotPolicy = Never, + stateCodec = Nothing + } + +sessionCodec :: Codec SessionEvent +sessionCodec = + Codec + { eventTypes = + EventType + <$> "SessionStarted" + :| [ "SessionCompleted", + "SessionFailed", + "SessionAwaiting", + "SessionResumed", + "InteractiveSessionRecorded", + "TurnRecorded" + ], + eventType = + EventType . \case + SessionStarted {} -> "SessionStarted" + SessionCompleted {} -> "SessionCompleted" + SessionFailed {} -> "SessionFailed" + SessionAwaiting {} -> "SessionAwaiting" + SessionResumed {} -> "SessionResumed" + InteractiveSessionRecorded {} -> "InteractiveSessionRecorded" + TurnRecorded {} -> "TurnRecorded", + schemaVersion = 1, + encode = toJSON, + decode = const parseSessionEvent, + upcasters = [] + } + +parseSessionEvent :: Value -> Either Text SessionEvent +parseSessionEvent value = + case parseEither parseJSON value of + Right event -> Right event + Left nativeErr -> + case parseEither parseLegacySessionEvent value of + Right event -> Right event + Left legacyErr -> Left (Text.pack nativeErr <> "; legacy decode failed: " <> Text.pack legacyErr) + +parseLegacySessionEvent :: Value -> Parser SessionEvent +parseLegacySessionEvent = + withObject "Rei AgentSessionEvent" $ \o -> do + tag <- o .: "type" + payload <- o .: "data" + case tag of + "agent_session_started" -> SessionStarted <$> parseLegacySessionStarted payload + "agent_session_completed" -> SessionCompleted <$> parseLegacySessionCompleted payload + "agent_session_failed" -> SessionFailed <$> parseLegacySessionFailed payload + "interactive_session_recorded" -> InteractiveSessionRecorded <$> parseLegacyInteractiveSessionRecorded payload + other -> fail ("Unknown Rei AgentSessionEvent tag: " <> Text.unpack other) + +parseLegacySessionStarted :: Value -> Parser SessionStartedData +parseLegacySessionStarted = + withObject "Rei AgentSessionStartedData" $ \o -> do + sessionId <- parseLegacySessionId =<< o .: "sessionId" + intentionId <- o .:? "intentionId" + previousSessionId <- traverse parseLegacySessionId =<< o .:? "previousSessionId" + SessionStartedData sessionId + <$> o .: "agentId" + <*> (normalizeLegacyFocus <$> o .: "focusType") + <*> pure (sessionScope intentionId) + <*> o .:? "focusTarget" + <*> pure previousSessionId + <*> pure Nothing + <*> pure 0 + <*> o .: "startedAt" + +parseLegacySessionCompleted :: Value -> Parser SessionCompletedData +parseLegacySessionCompleted = + withObject "Rei AgentSessionCompletedData" $ \o -> + SessionCompletedData + <$> (parseLegacySessionId =<< o .: "sessionId") + <*> o .: "completedAt" + <*> o .:? "modelUsed" + <*> o .:? "summary" + +parseLegacySessionFailed :: Value -> Parser SessionFailedData +parseLegacySessionFailed = + withObject "Rei AgentSessionFailedData" $ \o -> + SessionFailedData + <$> (parseLegacySessionId =<< o .: "sessionId") + <*> o .: "failedAt" + <*> o .: "errorMessage" + +parseLegacyInteractiveSessionRecorded :: Value -> Parser InteractiveSessionRecordedData +parseLegacyInteractiveSessionRecorded = + withObject "Rei InteractiveSessionRecordedData" $ \o -> do + sessionId <- parseLegacySessionId =<< o .: "sessionId" + intentionId <- o .:? "intentionId" + InteractiveSessionRecordedData sessionId + <$> o .: "agentId" + <*> (normalizeLegacyFocus <$> o .: "focusType") + <*> pure (sessionScope intentionId) + <*> pure Nothing + <*> o .: "startedAt" + +parseLegacySessionId :: Text -> Parser SessionId +parseLegacySessionId = either (fail . Text.unpack) pure . parseIdLenient + +sessionScope :: Maybe Text -> MemoryScope +sessionScope = \case + Just intentionId -> ScopeEntity reiNamespace (ScopeKind "intention") intentionId + Nothing -> ScopeGlobal reiNamespace + +reiNamespace :: Namespace +reiNamespace = Namespace "rei" + +normalizeLegacyFocus :: Text -> Text +normalizeLegacyFocus = \case + "FocusGeneralCoaching" -> "general_coaching" + "FocusToday" -> "today" + "FocusIntentionReview" -> "intention_review" + "FocusNudge" -> "nudge" + "FocusDailyReflection" -> "daily_reflection" + "FocusWeeklyReflection" -> "weekly_reflection" + "FocusNoteHelp" -> "note_help" + "FocusAssist" -> "assist" + "FocusIntentionAssist" -> "intention_assist" + "FocusCollectionExplore" -> "collection_explore" + "FocusCreateNote" -> "create_note" + "FocusCreateSkill" -> "create_skill" + "FocusScheduledWork" -> "scheduled_work" + "FocusUpdateNote" -> "update_note" + "FocusAskNote" -> "ask_note" + alreadyNormalized -> alreadyNormalized diff --git a/src/Kioku/Session/ReadModel.hs b/src/Kioku/Session/ReadModel.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Session/ReadModel.hs @@ -0,0 +1,670 @@ +module Kioku.Session.ReadModel + ( sessionInlineProjection, + SessionRow (..), + TurnRow (..), + SessionByIdQuery (..), + SessionsByNamespaceQuery (..), + SessionsByScopeQuery (..), + SessionsByFocusQuery (..), + SessionsByStartedRangeQuery (..), + SessionChainQuery (..), + SessionDelegationChildrenQuery (..), + AwaitingSessionsByCorrelationKeyQuery (..), + TurnsBySessionQuery (..), + sessionByIdReadModel, + sessionsByNamespaceReadModel, + sessionsByScopeReadModel, + sessionsByFocusReadModel, + sessionsByStartedRangeReadModel, + sessionChainReadModel, + sessionDelegationChildrenReadModel, + awaitingSessionsByCorrelationKeyReadModel, + turnsBySessionReadModel, + ) +where + +import Contravariant.Extras (contrazip2, contrazip3, contrazip4) +import Data.Functor.Contravariant ((>$<)) +import Data.Int (Int32) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Keiro.Projection (InlineProjection (..)) +import Keiro.ReadModel (ConsistencyMode (..), ReadModel (..), StrongScope (..)) +import Kioku.Api.Scope (scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Id (idText) +import Kioku.Prelude +import Kioku.Session.Domain +import Kiroku.Store.Types (RecordedEvent) + +data SessionRow = SessionRow + { sessionId :: !Text, + agentId :: !Text, + focus :: !Text, + namespace :: !Text, + scopeKind :: !(Maybe Text), + scopeRef :: !(Maybe Text), + subjectRef :: !(Maybe Text), + previousSessionId :: !(Maybe Text), + parentSessionId :: !(Maybe Text), + delegationDepth :: !Int, + status :: !Text, + startedAt :: !UTCTime, + completedAt :: !(Maybe UTCTime), + modelUsed :: !(Maybe Text), + summary :: !(Maybe Text), + errorMessage :: !(Maybe Text), + awaitingReason :: !(Maybe Text), + awaitingCorrelationKey :: !(Maybe Text), + awaitingDeadline :: !(Maybe UTCTime), + resumeInput :: !(Maybe Text) + } + deriving stock (Generic, Eq, Show) + +data TurnRow = TurnRow + { turnId :: !Text, + sessionId :: !Text, + turnIndex :: !Int, + role :: !Text, + content :: !Text, + toolSummary :: !(Maybe Text), + promptTokens :: !(Maybe Int), + outputTokens :: !(Maybe Int), + recordedAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +newtype SessionByIdQuery = SessionByIdQuery Text + +data SessionsByNamespaceQuery = SessionsByNamespaceQuery Text Int + +data SessionsByScopeQuery = SessionsByScopeQuery Text (Maybe Text) (Maybe Text) + +data SessionsByFocusQuery = SessionsByFocusQuery Text Text + +data SessionsByStartedRangeQuery = SessionsByStartedRangeQuery Text UTCTime UTCTime + +newtype SessionChainQuery = SessionChainQuery Text + +newtype SessionDelegationChildrenQuery = SessionDelegationChildrenQuery Text + +data AwaitingSessionsByCorrelationKeyQuery = AwaitingSessionsByCorrelationKeyQuery Text Text + +newtype TurnsBySessionQuery = TurnsBySessionQuery Text + +sessionInlineProjection :: InlineProjection SessionEvent +sessionInlineProjection = + InlineProjection + { name = "kioku-session-inline", + apply = applySessionEvent + } + +applySessionEvent :: SessionEvent -> RecordedEvent -> Tx.Transaction () +applySessionEvent event _recorded = + case event of + SessionStarted d -> Tx.statement (startedRow d) upsertSessionStmt + InteractiveSessionRecorded d -> Tx.statement (interactiveRow d) upsertSessionStmt + SessionCompleted d -> + Tx.statement + (idText d.sessionId, d.completedAt, d.modelUsed, d.summary) + updateSessionCompletedStmt + SessionFailed d -> + Tx.statement + (idText d.sessionId, d.failedAt, d.errorMessage) + updateSessionFailedStmt + SessionAwaiting d -> + Tx.statement + (idText d.sessionId, d.reason, d.correlationKey, d.deadline) + updateSessionAwaitingStmt + SessionResumed d -> + Tx.statement + (idText d.sessionId, d.input) + updateSessionResumedStmt + TurnRecorded d -> Tx.statement (turnRow d) insertTurnStmt + +startedRow :: SessionStartedData -> SessionRow +startedRow d = + SessionRow + { sessionId = idText d.sessionId, + agentId = d.agentId, + focus = d.focus, + namespace = scopeNamespaceText d.scope, + scopeKind = scopeKindText d.scope, + scopeRef = scopeRefText d.scope, + subjectRef = d.subjectRef, + previousSessionId = idText <$> d.previousSessionId, + parentSessionId = idText <$> d.parentSessionId, + delegationDepth = d.delegationDepth, + status = "running", + startedAt = d.startedAt, + completedAt = Nothing, + modelUsed = Nothing, + summary = Nothing, + errorMessage = Nothing, + awaitingReason = Nothing, + awaitingCorrelationKey = Nothing, + awaitingDeadline = Nothing, + resumeInput = Nothing + } + +interactiveRow :: InteractiveSessionRecordedData -> SessionRow +interactiveRow d = + SessionRow + { sessionId = idText d.sessionId, + agentId = d.agentId, + focus = d.focus, + namespace = scopeNamespaceText d.scope, + scopeKind = scopeKindText d.scope, + scopeRef = scopeRefText d.scope, + subjectRef = d.subjectRef, + previousSessionId = Nothing, + parentSessionId = Nothing, + delegationDepth = 0, + status = "interactive", + startedAt = d.startedAt, + completedAt = Nothing, + modelUsed = Nothing, + summary = Nothing, + errorMessage = Nothing, + awaitingReason = Nothing, + awaitingCorrelationKey = Nothing, + awaitingDeadline = Nothing, + resumeInput = Nothing + } + +turnRow :: TurnRecordedData -> TurnRow +turnRow d = + TurnRow + { turnId = d.turnId, + sessionId = idText d.sessionId, + turnIndex = d.turnIndex, + role = d.role, + content = d.content, + toolSummary = d.toolSummary, + promptTokens = d.promptTokens, + outputTokens = d.outputTokens, + recordedAt = d.recordedAt + } + +sessionByIdReadModel :: ReadModel SessionByIdQuery (Maybe SessionRow) +sessionByIdReadModel = + ReadModel + { name = "kioku-session-by-id", + schema = "kiroku", + tableName = "kioku_sessions", + subscriptionName = "kioku-session-inline", + version = 3, + shapeHash = "kioku-session-v3", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(SessionByIdQuery sid) -> Tx.statement sid selectSessionByIdStmt + } + +sessionsByNamespaceReadModel :: ReadModel SessionsByNamespaceQuery [SessionRow] +sessionsByNamespaceReadModel = + ReadModel + { name = "kioku-sessions-by-namespace", + schema = "kiroku", + tableName = "kioku_sessions", + subscriptionName = "kioku-session-inline", + version = 3, + shapeHash = "kioku-session-v3", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \q -> Tx.statement q selectSessionsByNamespaceStmt + } + +sessionsByScopeReadModel :: ReadModel SessionsByScopeQuery [SessionRow] +sessionsByScopeReadModel = + ReadModel + { name = "kioku-sessions-by-scope", + schema = "kiroku", + tableName = "kioku_sessions", + subscriptionName = "kioku-session-inline", + version = 3, + shapeHash = "kioku-session-v3", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \q -> Tx.statement q selectSessionsByScopeStmt + } + +sessionsByFocusReadModel :: ReadModel SessionsByFocusQuery [SessionRow] +sessionsByFocusReadModel = + ReadModel + { name = "kioku-sessions-by-focus", + schema = "kiroku", + tableName = "kioku_sessions", + subscriptionName = "kioku-session-inline", + version = 3, + shapeHash = "kioku-session-v3", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \q -> Tx.statement q selectSessionsByFocusStmt + } + +sessionsByStartedRangeReadModel :: ReadModel SessionsByStartedRangeQuery [SessionRow] +sessionsByStartedRangeReadModel = + ReadModel + { name = "kioku-sessions-by-started-range", + schema = "kiroku", + tableName = "kioku_sessions", + subscriptionName = "kioku-session-inline", + version = 3, + shapeHash = "kioku-session-v3", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \q -> Tx.statement q selectSessionsByStartedRangeStmt + } + +sessionChainReadModel :: ReadModel SessionChainQuery [SessionRow] +sessionChainReadModel = + ReadModel + { name = "kioku-session-chain", + schema = "kiroku", + tableName = "kioku_sessions", + subscriptionName = "kioku-session-inline", + version = 3, + shapeHash = "kioku-session-v3", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(SessionChainQuery sid) -> Tx.statement sid selectSessionChainStmt + } + +sessionDelegationChildrenReadModel :: ReadModel SessionDelegationChildrenQuery [SessionRow] +sessionDelegationChildrenReadModel = + ReadModel + { name = "kioku-session-delegation-children", + schema = "kiroku", + tableName = "kioku_sessions", + subscriptionName = "kioku-session-inline", + version = 3, + shapeHash = "kioku-session-v3", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(SessionDelegationChildrenQuery sid) -> Tx.statement sid selectDelegationChildrenStmt + } + +awaitingSessionsByCorrelationKeyReadModel :: ReadModel AwaitingSessionsByCorrelationKeyQuery [SessionRow] +awaitingSessionsByCorrelationKeyReadModel = + ReadModel + { name = "kioku-sessions-awaiting-by-correlation-key", + schema = "kiroku", + tableName = "kioku_sessions", + subscriptionName = "kioku-session-inline", + version = 3, + shapeHash = "kioku-session-v3", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \q -> Tx.statement q selectAwaitingByCorrelationKeyStmt + } + +turnsBySessionReadModel :: ReadModel TurnsBySessionQuery [TurnRow] +turnsBySessionReadModel = + ReadModel + { name = "kioku-turns-by-session", + schema = "kiroku", + tableName = "kioku_turns", + subscriptionName = "kioku-session-inline", + version = 1, + shapeHash = "kioku-turn-v1", + defaultConsistency = Eventual, + strongScope = EntireLog, + query = \(TurnsBySessionQuery sid) -> Tx.statement sid selectTurnsBySessionStmt + } + +sessionRowDecoder :: D.Row SessionRow +sessionRowDecoder = + SessionRow + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> (fromIntegral @Int32 @Int <$> D.column (D.nonNullable D.int4)) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.timestamptz) + <*> D.column (D.nullable D.timestamptz) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.timestamptz) + <*> D.column (D.nullable D.text) + +turnRowDecoder :: D.Row TurnRow +turnRowDecoder = + TurnRow + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> (fromIntegral @Int32 @Int <$> D.column (D.nonNullable D.int4)) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + <*> (fmap (fromIntegral @Int32 @Int) <$> D.column (D.nullable D.int4)) + <*> (fmap (fromIntegral @Int32 @Int) <$> D.column (D.nullable D.int4)) + <*> D.column (D.nonNullable D.timestamptz) + +selectSessionByIdStmt :: Statement Text (Maybe SessionRow) +selectSessionByIdStmt = + preparable + """ + SELECT session_id, agent_id, focus, namespace, scope_kind, scope_ref, subject_ref, + previous_session_id, parent_session_id, delegation_depth, status, started_at, + completed_at, model_used, summary, error_message, + awaiting_reason, awaiting_correlation_key, awaiting_deadline, resume_input + FROM kioku_sessions + WHERE session_id = $1 + """ + (E.param (E.nonNullable E.text)) + (D.rowMaybe sessionRowDecoder) + +selectSessionsByNamespaceStmt :: Statement SessionsByNamespaceQuery [SessionRow] +selectSessionsByNamespaceStmt = + preparable + """ + SELECT session_id, agent_id, focus, namespace, scope_kind, scope_ref, subject_ref, + previous_session_id, parent_session_id, delegation_depth, status, started_at, + completed_at, model_used, summary, error_message, + awaiting_reason, awaiting_correlation_key, awaiting_deadline, resume_input + FROM kioku_sessions + WHERE namespace = $1 + ORDER BY started_at DESC + LIMIT $2 + """ + ( ((\(SessionsByNamespaceQuery ns _) -> ns) >$< E.param (E.nonNullable E.text)) + <> ((\(SessionsByNamespaceQuery _ limit) -> fromIntegral @Int @Int32 limit) >$< E.param (E.nonNullable E.int4)) + ) + (D.rowList sessionRowDecoder) + +selectSessionsByScopeStmt :: Statement SessionsByScopeQuery [SessionRow] +selectSessionsByScopeStmt = + preparable + """ + SELECT session_id, agent_id, focus, namespace, scope_kind, scope_ref, subject_ref, + previous_session_id, parent_session_id, delegation_depth, status, started_at, + completed_at, model_used, summary, error_message, + awaiting_reason, awaiting_correlation_key, awaiting_deadline, resume_input + FROM kioku_sessions + WHERE namespace = $1 + AND scope_kind IS NOT DISTINCT FROM $2 + AND scope_ref IS NOT DISTINCT FROM $3 + ORDER BY started_at DESC + """ + ( ((\(SessionsByScopeQuery ns _ _) -> ns) >$< E.param (E.nonNullable E.text)) + <> ((\(SessionsByScopeQuery _ sk _) -> sk) >$< E.param (E.nullable E.text)) + <> ((\(SessionsByScopeQuery _ _ sr) -> sr) >$< E.param (E.nullable E.text)) + ) + (D.rowList sessionRowDecoder) + +selectSessionsByFocusStmt :: Statement SessionsByFocusQuery [SessionRow] +selectSessionsByFocusStmt = + preparable + """ + SELECT session_id, agent_id, focus, namespace, scope_kind, scope_ref, subject_ref, + previous_session_id, parent_session_id, delegation_depth, status, started_at, + completed_at, model_used, summary, error_message, + awaiting_reason, awaiting_correlation_key, awaiting_deadline, resume_input + FROM kioku_sessions + WHERE namespace = $1 + AND focus = $2 + ORDER BY started_at DESC + """ + ( ((\(SessionsByFocusQuery ns _) -> ns) >$< E.param (E.nonNullable E.text)) + <> ((\(SessionsByFocusQuery _ focus) -> focus) >$< E.param (E.nonNullable E.text)) + ) + (D.rowList sessionRowDecoder) + +selectSessionsByStartedRangeStmt :: Statement SessionsByStartedRangeQuery [SessionRow] +selectSessionsByStartedRangeStmt = + preparable + """ + SELECT session_id, agent_id, focus, namespace, scope_kind, scope_ref, subject_ref, + previous_session_id, parent_session_id, delegation_depth, status, started_at, + completed_at, model_used, summary, error_message, + awaiting_reason, awaiting_correlation_key, awaiting_deadline, resume_input + FROM kioku_sessions + WHERE namespace = $1 + AND started_at >= $2 + AND started_at < $3 + ORDER BY started_at DESC + """ + ( ((\(SessionsByStartedRangeQuery ns _ _) -> ns) >$< E.param (E.nonNullable E.text)) + <> ((\(SessionsByStartedRangeQuery _ start _) -> start) >$< E.param (E.nonNullable E.timestamptz)) + <> ((\(SessionsByStartedRangeQuery _ _ end) -> end) >$< E.param (E.nonNullable E.timestamptz)) + ) + (D.rowList sessionRowDecoder) + +-- | Walk a session's continuation chain backwards through @previous_session_id@. +-- +-- The @path@ array makes revisiting a session impossible, so the walk terminates on any +-- data — including a cycle, which 'Kioku.Session.start' now refuses to create but which +-- unchecked lineage (see @validateLineage@) still permits to be constructed out of order. +-- With a plain @UNION ALL@ and no guard, a cycle loops until timeout or OOM. The depth cap +-- is defense in depth, far above any legitimate chain. +-- +-- The final @SELECT@ omits @depth@/@path@, so 'sessionRowDecoder' is unchanged. +selectSessionChainStmt :: Statement Text [SessionRow] +selectSessionChainStmt = + preparable + """ + WITH RECURSIVE chain AS ( + SELECT session_id, agent_id, focus, namespace, scope_kind, scope_ref, subject_ref, + previous_session_id, parent_session_id, delegation_depth, status, started_at, + completed_at, model_used, summary, error_message, + awaiting_reason, awaiting_correlation_key, awaiting_deadline, resume_input, + 1 AS depth, ARRAY[session_id] AS path + FROM kioku_sessions + WHERE session_id = $1 + UNION ALL + SELECT s.session_id, s.agent_id, s.focus, s.namespace, s.scope_kind, s.scope_ref, s.subject_ref, + s.previous_session_id, s.parent_session_id, s.delegation_depth, s.status, s.started_at, + s.completed_at, s.model_used, s.summary, s.error_message, + s.awaiting_reason, s.awaiting_correlation_key, s.awaiting_deadline, s.resume_input, + c.depth + 1, c.path || s.session_id + FROM kioku_sessions s + INNER JOIN chain c ON s.session_id = c.previous_session_id + WHERE NOT s.session_id = ANY (c.path) + AND c.depth < 10000 + ) + SELECT session_id, agent_id, focus, namespace, scope_kind, scope_ref, subject_ref, + previous_session_id, parent_session_id, delegation_depth, status, started_at, + completed_at, model_used, summary, error_message, + awaiting_reason, awaiting_correlation_key, awaiting_deadline, resume_input + FROM chain + ORDER BY started_at ASC + """ + (E.param (E.nonNullable E.text)) + (D.rowList sessionRowDecoder) + +selectDelegationChildrenStmt :: Statement Text [SessionRow] +selectDelegationChildrenStmt = + preparable + """ + SELECT session_id, agent_id, focus, namespace, scope_kind, scope_ref, subject_ref, + previous_session_id, parent_session_id, delegation_depth, status, started_at, + completed_at, model_used, summary, error_message, + awaiting_reason, awaiting_correlation_key, awaiting_deadline, resume_input + FROM kioku_sessions + WHERE parent_session_id = $1 + ORDER BY started_at ASC, session_id ASC + """ + (E.param (E.nonNullable E.text)) + (D.rowList sessionRowDecoder) + +selectTurnsBySessionStmt :: Statement Text [TurnRow] +selectTurnsBySessionStmt = + preparable + """ + SELECT turn_id, session_id, turn_index, role, content, tool_summary, prompt_tokens, + output_tokens, recorded_at + FROM kioku_turns + WHERE session_id = $1 + ORDER BY turn_index ASC + """ + (E.param (E.nonNullable E.text)) + (D.rowList turnRowDecoder) + +upsertSessionStmt :: Statement SessionRow () +upsertSessionStmt = + preparable + """ + INSERT INTO kioku_sessions + (session_id, agent_id, focus, namespace, scope_kind, scope_ref, subject_ref, + previous_session_id, parent_session_id, delegation_depth, status, started_at, + completed_at, model_used, summary, error_message, + awaiting_reason, awaiting_correlation_key, awaiting_deadline, resume_input, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, NOW()) + ON CONFLICT (session_id) DO UPDATE SET + agent_id = EXCLUDED.agent_id, + focus = EXCLUDED.focus, + namespace = EXCLUDED.namespace, + scope_kind = EXCLUDED.scope_kind, + scope_ref = EXCLUDED.scope_ref, + subject_ref = EXCLUDED.subject_ref, + previous_session_id = EXCLUDED.previous_session_id, + parent_session_id = EXCLUDED.parent_session_id, + delegation_depth = EXCLUDED.delegation_depth, + status = EXCLUDED.status, + started_at = EXCLUDED.started_at, + completed_at = EXCLUDED.completed_at, + model_used = EXCLUDED.model_used, + summary = EXCLUDED.summary, + error_message = EXCLUDED.error_message, + awaiting_reason = EXCLUDED.awaiting_reason, + awaiting_correlation_key = EXCLUDED.awaiting_correlation_key, + awaiting_deadline = EXCLUDED.awaiting_deadline, + resume_input = EXCLUDED.resume_input, + updated_at = EXCLUDED.updated_at + """ + sessionRowEncoder + D.noResult + +sessionRowEncoder :: E.Params SessionRow +sessionRowEncoder = + ((\row -> row.sessionId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.agentId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.focus) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.namespace) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.scopeKind) >$< E.param (E.nullable E.text)) + <> ((\row -> row.scopeRef) >$< E.param (E.nullable E.text)) + <> ((\row -> row.subjectRef) >$< E.param (E.nullable E.text)) + <> ((\row -> row.previousSessionId) >$< E.param (E.nullable E.text)) + <> ((\row -> row.parentSessionId) >$< E.param (E.nullable E.text)) + <> ((\row -> fromIntegral @Int @Int32 row.delegationDepth) >$< E.param (E.nonNullable E.int4)) + <> ((\row -> row.status) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.startedAt) >$< E.param (E.nonNullable E.timestamptz)) + <> ((\row -> row.completedAt) >$< E.param (E.nullable E.timestamptz)) + <> ((\row -> row.modelUsed) >$< E.param (E.nullable E.text)) + <> ((\row -> row.summary) >$< E.param (E.nullable E.text)) + <> ((\row -> row.errorMessage) >$< E.param (E.nullable E.text)) + <> ((\row -> row.awaitingReason) >$< E.param (E.nullable E.text)) + <> ((\row -> row.awaitingCorrelationKey) >$< E.param (E.nullable E.text)) + <> ((\row -> row.awaitingDeadline) >$< E.param (E.nullable E.timestamptz)) + <> ((\row -> row.resumeInput) >$< E.param (E.nullable E.text)) + +updateSessionCompletedStmt :: Statement (Text, UTCTime, Maybe Text, Maybe Text) () +updateSessionCompletedStmt = + preparable + "UPDATE kioku_sessions SET status = 'completed', completed_at = $2, model_used = $3, summary = $4, awaiting_reason = NULL, awaiting_correlation_key = NULL, awaiting_deadline = NULL, updated_at = NOW() WHERE session_id = $1" + ( contrazip4 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.timestamptz)) + (E.param (E.nullable E.text)) + (E.param (E.nullable E.text)) + ) + D.noResult + +updateSessionFailedStmt :: Statement (Text, UTCTime, Text) () +updateSessionFailedStmt = + preparable + "UPDATE kioku_sessions SET status = 'failed', completed_at = $2, error_message = $3, awaiting_reason = NULL, awaiting_correlation_key = NULL, awaiting_deadline = NULL, updated_at = NOW() WHERE session_id = $1" + ( contrazip3 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.timestamptz)) + (E.param (E.nonNullable E.text)) + ) + D.noResult + +-- | Park the session. @resume_input@ is cleared so a re-park does not leave the previous +-- wait's answer visible on the row. +-- +-- @awaiting_deadline@ is advisory only: it is stored for hosts, and kioku does not enforce +-- it (no timer fires, nothing expires). See MasterPlan 2's Decision Log (2026-07-07). +updateSessionAwaitingStmt :: Statement (Text, Text, Maybe Text, Maybe UTCTime) () +updateSessionAwaitingStmt = + preparable + "UPDATE kioku_sessions SET status = 'awaiting', awaiting_reason = $2, awaiting_correlation_key = $3, awaiting_deadline = $4, resume_input = NULL, updated_at = NOW() WHERE session_id = $1" + ( contrazip4 + (E.param (E.nonNullable E.text)) + (E.param (E.nonNullable E.text)) + (E.param (E.nullable E.text)) + (E.param (E.nullable E.timestamptz)) + ) + D.noResult + +updateSessionResumedStmt :: Statement (Text, Text) () +updateSessionResumedStmt = + preparable + "UPDATE kioku_sessions SET status = 'running', resume_input = $2, awaiting_reason = NULL, awaiting_correlation_key = NULL, awaiting_deadline = NULL, updated_at = NOW() WHERE session_id = $1" + (contrazip2 (E.param (E.nonNullable E.text)) (E.param (E.nonNullable E.text))) + D.noResult + +selectAwaitingByCorrelationKeyStmt :: Statement AwaitingSessionsByCorrelationKeyQuery [SessionRow] +selectAwaitingByCorrelationKeyStmt = + preparable + """ + SELECT session_id, agent_id, focus, namespace, scope_kind, scope_ref, subject_ref, + previous_session_id, parent_session_id, delegation_depth, status, started_at, + completed_at, model_used, summary, error_message, + awaiting_reason, awaiting_correlation_key, awaiting_deadline, resume_input + FROM kioku_sessions + WHERE namespace = $1 + AND status = 'awaiting' + AND awaiting_correlation_key = $2 + ORDER BY started_at DESC + """ + ( ((\(AwaitingSessionsByCorrelationKeyQuery ns _) -> ns) >$< E.param (E.nonNullable E.text)) + <> ((\(AwaitingSessionsByCorrelationKeyQuery _ k) -> k) >$< E.param (E.nonNullable E.text)) + ) + (D.rowList sessionRowDecoder) + +-- | @(session_id, turn_index)@ is the turn's identity; @turn_id@ is an idempotency token +-- that travels with it. The conflict clause updates @turn_id@ as well, so rebuilding the +-- projection from the event stream cannot leave a superseded turn's id attached to the +-- winning event's content. +insertTurnStmt :: Statement TurnRow () +insertTurnStmt = + preparable + """ + INSERT INTO kioku_turns + (turn_id, session_id, turn_index, role, content, tool_summary, prompt_tokens, output_tokens, recorded_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) + ON CONFLICT (session_id, turn_index) DO UPDATE SET + turn_id = EXCLUDED.turn_id, + role = EXCLUDED.role, + content = EXCLUDED.content, + tool_summary = EXCLUDED.tool_summary, + prompt_tokens = EXCLUDED.prompt_tokens, + output_tokens = EXCLUDED.output_tokens, + recorded_at = EXCLUDED.recorded_at + """ + turnRowEncoder + D.noResult + +turnRowEncoder :: E.Params TurnRow +turnRowEncoder = + ((\row -> row.turnId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.sessionId) >$< E.param (E.nonNullable E.text)) + <> ((\row -> fromIntegral @Int @Int32 row.turnIndex) >$< E.param (E.nonNullable E.int4)) + <> ((\row -> row.role) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.content) >$< E.param (E.nonNullable E.text)) + <> ((\row -> row.toolSummary) >$< E.param (E.nullable E.text)) + <> ((\row -> fmap (fromIntegral @Int @Int32) row.promptTokens) >$< E.param (E.nullable E.int4)) + <> ((\row -> fmap (fromIntegral @Int @Int32) row.outputTokens) >$< E.param (E.nullable E.int4)) + <> ((\row -> row.recordedAt) >$< E.param (E.nonNullable E.timestamptz)) diff --git a/src/Kioku/Worker/Failure.hs b/src/Kioku/Worker/Failure.hs new file mode 100644 --- /dev/null +++ b/src/Kioku/Worker/Failure.hs @@ -0,0 +1,60 @@ +-- | Failure classification shared by kioku's background workers. +-- +-- Both workers must answer the same question about every failure: is this +-- worth retrying, is it permanently broken, or does it not belong to me? This +-- module owns the half of that taxonomy that concerns kiroku store errors and +-- the embedding worker's retry schedule. The distillation timers' half lives in +-- "Kioku.Distill.Timer.Outcome". +module Kioku.Worker.Failure + ( isTransientStoreError, + embeddingRetryDelay, + ) +where + +import Kiroku.Store.Error (StoreError (..)) +import Shibuya.Core.Ack (RetryDelay (..)) +import Shibuya.Core.Types (Attempt (..)) + +-- | Is this store error worth retrying? +-- +-- Transient errors are the ones kiroku's own constructor documentation calls +-- retryable: a pool timeout, a connection torn down mid-operation, and the +-- catch-all connection error. Everything else describes a stable disagreement +-- between the caller and the database — a version conflict, a missing stream, a +-- server error whose SQLSTATE kiroku does not recognise — and would fail +-- identically on every redelivery. +-- +-- The match lists every constructor explicitly rather than falling through a +-- wildcard so that @-Wincomplete-patterns@ forces a classification decision if +-- a future kiroku version adds one. +isTransientStoreError :: StoreError -> Bool +isTransientStoreError = \case + PoolAcquisitionTimeout -> True + ConnectionLost _ -> True + ConnectionError _ -> True + WrongExpectedVersion {} -> False + EmptyAppendBatch _ -> False + StreamNotFound _ -> False + ReservedStreamName _ -> False + StreamNameTooLong _ _ -> False + StreamAlreadyExists _ -> False + DuplicateEvent _ -> False + EventAlreadyLinked _ _ -> False + LinkSourceEventMissing _ -> False + UnexpectedServerError _ _ -> False + +-- | Backoff for a redelivered @MemoryRecorded@ event, by zero-based delivery +-- attempt: 5s, 20s, 60s, then 180s. +-- +-- The bound lives upstream, not here: the shibuya-kiroku adapter maps 'AckRetry' +-- onto kiroku's per-subscription retry policy, which dead-letters after five +-- total deliveries. At most four of these delays are ever consumed, and an +-- outage outlasting that window is recovered by the worker's startup backfill +-- rather than by retrying forever. 'Nothing' (an adapter that does not track +-- redeliveries) gets the longest delay. +embeddingRetryDelay :: Maybe Attempt -> RetryDelay +embeddingRetryDelay = \case + Just (Attempt 0) -> RetryDelay 5 + Just (Attempt 1) -> RetryDelay 20 + Just (Attempt 2) -> RetryDelay 60 + _ -> RetryDelay 180 diff --git a/test/Kioku/AwaitingSpec.hs b/test/Kioku/AwaitingSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/AwaitingSpec.hs @@ -0,0 +1,301 @@ +module Kioku.AwaitingSpec (tests) where + +import Data.Vector qualified as Vector +import Effectful (Eff, IOE, (:>)) +import Effectful.Error.Static (Error) +import Keiro.Stream qualified as Stream +import Kioku.Api.Scope (MemoryScope (..), Namespace (..)) +import Kioku.App (AppEffects, runAppIO, withNoopAppEnv) +import Kioku.Id (SessionId, genSessionId, idText) +import Kioku.Migrations.TestSupport (withKiokuMigratedDatabase) +import Kioku.Prelude +import Kioku.Session qualified as Session +import Kioku.Session.Domain + ( AwaitInputData (..), + CompleteSessionData (..), + FailSessionData (..), + ResumeSessionData (..), + SessionEvent (..), + StartSessionData (..), + ) +import Kioku.Session.EventStream (parseSessionEvent, sessionStream) +import Kioku.Session.ReadModel (SessionRow (..)) +import Kiroku.Store.Connection (defaultConnectionSettings) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Read (readStreamForward) +import Kiroku.Store.Types (RecordedEvent (..), StreamVersion (..)) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, assertFailure, testCase) + +tests :: TestTree +tests = + testGroup + "Awaiting park-and-resume" + [ testCase "park then resume" testParkAndResume, + testCase "find awaiting by correlation key" testFindAwaitingByCorrelationKey, + testCase "reconstruct aggregate after crash" testReconstructAfterCrash, + testCase "idempotent resume on re-delivery" testIdempotentResume, + testCase "correlation mismatch is rejected" testCorrelationMismatchRejected, + testCase "complete an awaiting session" testCompleteAwaitingSession, + testCase "fail an awaiting session" testFailAwaitingSession + ] + +testParkAndResume :: Assertion +testParkAndResume = + withAwaitingApp do + sid <- startFixture + parkFixture sid "approval_req_1" + parked <- getExisting sid + liftIO do + assertEqual "parked status" "awaiting" parked.status + assertEqual "awaiting reason" (Just "approval") parked.awaitingReason + assertEqual "awaiting correlation key" (Just "approval_req_1") parked.awaitingCorrelationKey + eventsAfterPark <- readSessionEvents sid + liftIO $ + assertBool "tail event is SessionAwaiting" $ + case lastMay eventsAfterPark of + Just SessionAwaiting {} -> True + _ -> False + resumeFixture sid (Just "approval_req_1") "approved" + resumed <- getExisting sid + liftIO do + assertEqual "resumed status" "running" resumed.status + assertEqual "resume input" (Just "approved") resumed.resumeInput + assertEqual "awaiting reason cleared" Nothing resumed.awaitingReason + assertEqual "awaiting key cleared" Nothing resumed.awaitingCorrelationKey + assertEqual "awaiting deadline cleared" Nothing resumed.awaitingDeadline + eventsAfterResume <- readSessionEvents sid + liftIO $ + assertBool "SessionResumed was appended" $ + any isSessionResumed eventsAfterResume + +testFindAwaitingByCorrelationKey :: Assertion +testFindAwaitingByCorrelationKey = + withAwaitingApp do + sid1 <- startFixture + sid2 <- startFixture + parkFixture sid1 "approval_req_1" + parkFixture sid2 "approval_req_2" + found <- Session.getAwaitingByCorrelationKey testNamespace "approval_req_1" >>= liftEither "getAwaitingByCorrelationKey" + liftIO $ + assertEqual "only matching parked session is returned" [idText sid1] (map (.sessionId) found) + +testReconstructAfterCrash :: Assertion +testReconstructAfterCrash = + withAwaitingApp do + sid <- startFixture + parkFixture sid "approval_req_1" + eventsAfterPark <- readSessionEvents sid + liftIO $ + assertEqual "events end parked" ["SessionStarted", "SessionAwaiting"] (eventName <$> eventsAfterPark) + resumeFixture sid (Just "approval_req_1") "approved" + resumed <- getExisting sid + liftIO $ + assertEqual "resume after log-backed park succeeds" "running" resumed.status + +testIdempotentResume :: Assertion +testIdempotentResume = + withAwaitingApp do + sid <- startFixture + parkFixture sid "approval_req_1" + resumeFixture sid (Just "approval_req_1") "approved" + resumeFixture sid (Just "approval_req_1") "approved" + eventsAfterResume <- readSessionEvents sid + liftIO $ + assertEqual "only one SessionResumed event is emitted" 1 (length (filter isSessionResumed eventsAfterResume)) + +testCorrelationMismatchRejected :: Assertion +testCorrelationMismatchRejected = + withAwaitingApp do + sid <- startFixture + parkFixture sid "k1" + now <- liftIO getCurrentTime + result <- + Session.resume + ResumeSessionData + { sessionId = sid, + correlationKey = Just "k2", + force = False, + input = "approved", + resumedAt = now + } + case result of + Left Session.SessionCorrelationMismatch -> pure () + other -> liftIO (assertFailure ("expected SessionCorrelationMismatch, got " <> show other)) + eventsAfterRejectedResume <- readSessionEvents sid + liftIO $ + assertBool "mismatched resume emits no SessionResumed event" $ + not (any isSessionResumed eventsAfterRejectedResume) + +testCompleteAwaitingSession :: Assertion +testCompleteAwaitingSession = + withAwaitingApp do + sid <- startFixture + parkFixture sid "approval_req_1" + now <- liftIO getCurrentTime + completeResult <- + Session.complete + CompleteSessionData + { sessionId = sid, + completedAt = now, + modelUsed = Just "test-model", + summary = Just "completed while parked" + } + void (liftEither "Session.complete" completeResult) + completed <- getExisting sid + events <- readSessionEvents sid + liftIO do + assertEqual "completed status" "completed" completed.status + assertEqual "awaiting reason cleared" Nothing completed.awaitingReason + assertEqual "awaiting key cleared" Nothing completed.awaitingCorrelationKey + assertEqual "awaiting deadline cleared" Nothing completed.awaitingDeadline + assertEqual "completed stream" ["SessionStarted", "SessionAwaiting", "SessionCompleted"] (eventName <$> events) + +testFailAwaitingSession :: Assertion +testFailAwaitingSession = + withAwaitingApp do + sid <- startFixture + parkFixture sid "approval_req_1" + now <- liftIO getCurrentTime + failResult <- + Session.failSession + FailSessionData + { sessionId = sid, + failedAt = now, + errorMessage = "timed out" + } + void (liftEither "Session.failSession" failResult) + failed <- getExisting sid + events <- readSessionEvents sid + liftIO do + assertEqual "failed status" "failed" failed.status + assertEqual "error message" (Just "timed out") failed.errorMessage + assertEqual "awaiting reason cleared" Nothing failed.awaitingReason + assertEqual "awaiting key cleared" Nothing failed.awaitingCorrelationKey + assertEqual "awaiting deadline cleared" Nothing failed.awaitingDeadline + assertEqual "failed stream" ["SessionStarted", "SessionAwaiting", "SessionFailed"] (eventName <$> events) + +withAwaitingApp :: + Eff AppEffects a -> + IO a +withAwaitingApp action = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) \env -> do + result <- runAppIO env action + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right value -> pure value + +startFixture :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + Eff es SessionId +startFixture = do + sid <- liftIO genSessionId + now <- liftIO getCurrentTime + result <- + Session.start + StartSessionData + { sessionId = sid, + agentId = "test-agent", + focus = "awaiting lifecycle", + scope = testScope, + subjectRef = Nothing, + previousSessionId = Nothing, + parentSessionId = Nothing, + delegationDepth = 0, + startedAt = now + } + void (liftEither "Session.start" result) + pure sid + +parkFixture :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + Text -> + Eff es () +parkFixture sid key = do + now <- liftIO getCurrentTime + result <- + Session.awaitInput + AwaitInputData + { sessionId = sid, + reason = "approval", + correlationKey = Just key, + deadline = Nothing, + awaitedAt = now + } + void (liftEither "Session.awaitInput" result) + +resumeFixture :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + Maybe Text -> + Text -> + Eff es () +resumeFixture sid key input = do + now <- liftIO getCurrentTime + result <- + Session.resume + ResumeSessionData + { sessionId = sid, + correlationKey = key, + force = False, + input, + resumedAt = now + } + void (liftEither "Session.resume" result) + +getExisting :: + (IOE :> es, Store :> es) => + SessionId -> + Eff es SessionRow +getExisting sid = do + result <- Session.getById sid >>= liftEither "Session.getById" + case result of + Nothing -> liftIO (assertFailure ("missing session row " <> show (idText sid))) + Just row -> pure row + +readSessionEvents :: + (IOE :> es, Store :> es) => + SessionId -> + Eff es [SessionEvent] +readSessionEvents sid = do + recorded <- Vector.toList <$> readStreamForward (Stream.streamName (sessionStream sid)) (StreamVersion 0) 100 + traverse decodeRecorded recorded + where + decodeRecorded recorded = + case parseSessionEvent recorded.payload of + Left err -> liftIO (assertFailure ("parseSessionEvent: " <> show err)) + Right event -> pure event + +liftEither :: (Show e, IOE :> es) => String -> Either e a -> Eff es a +liftEither label = \case + Left err -> liftIO (assertFailure (label <> ": " <> show err)) + Right value -> pure value + +isSessionResumed :: SessionEvent -> Bool +isSessionResumed = \case + SessionResumed {} -> True + _ -> False + +eventName :: SessionEvent -> Text +eventName = \case + SessionStarted {} -> "SessionStarted" + SessionAwaiting {} -> "SessionAwaiting" + SessionResumed {} -> "SessionResumed" + SessionCompleted {} -> "SessionCompleted" + SessionFailed {} -> "SessionFailed" + InteractiveSessionRecorded {} -> "InteractiveSessionRecorded" + TurnRecorded {} -> "TurnRecorded" + +lastMay :: [a] -> Maybe a +lastMay [] = Nothing +lastMay xs = Just (last xs) + +testNamespace :: Namespace +testNamespace = Namespace "kioku-test" + +testScope :: MemoryScope +testScope = ScopeGlobal testNamespace diff --git a/test/Kioku/DistillSpec.hs b/test/Kioku/DistillSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/DistillSpec.hs @@ -0,0 +1,1614 @@ +{-# LANGUAGE DataKinds #-} + +module Kioku.DistillSpec + ( tests, + ) +where + +import Baikai (Response, _Response, _TextContent) +import Baikai.Content (AssistantContent (..)) +import Baikai.Embedding (EmbeddingModel) +import Data.Aeson qualified as Aeson +import Data.ByteString.Lazy qualified as BL +import Data.Foldable (traverse_) +import Data.Functor.Contravariant ((>$<)) +import Data.IORef (IORef, modifyIORef', newIORef, readIORef) +import Data.Int (Int64) +import Data.Set qualified as Set +import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE +import Data.Text.IO qualified as TextIO +import Data.Time (addUTCTime, diffUTCTime) +import Data.Vector qualified as Vector +import Effectful (Eff, IOE, runEff, (:>)) +import Effectful.Dispatch.Dynamic (interpret) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (runPrim) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Keiro.Stream qualified as Stream +import Keiro.Timer (countDueTimers) +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), ScopeKind (..), scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Api.Types (Confidence (..), MemoryType (..)) +import Kioku.App (AppEnv, runAppIO, withNoopAppEnv) +import Kioku.Distill.Consolidate (ConsolidateInput (..), ConsolidationAction (..), ConsolidationDecision (..), ExistingMemory (..), consolidateProgram) +import Kioku.Distill.Extract (ExtractOutput (..), ExtractedAtom (..), extractProgram) +import Kioku.Distill.L1 (L1Error (..), L1Outcome (..), L1RunMode (..), L1Summary (..), distillSessionL1, recallCandidates, scopedScanCandidates) +import Kioku.Distill.L2 (SceneRow (..), getScenesByScope, regenerateScene, sceneMirrorPath) +import Kioku.Distill.L3 (PersonaRow (..), getPersonaByScope, personaMirrorPath, regeneratePersona) +import Kioku.Distill.Persona (personaProgram) +import Kioku.Distill.Runtime (DistillRuntime (..), newDistillRuntime) +import Kioku.Distill.Scene (SceneInput (..), sceneProgram) +import Kioku.Distill.Timer (idleFlushSeconds, l1ExtractProcessManagerName) +import Kioku.Distill.Timer.Worker (runKiokuTimerWorkerOnce) +import Kioku.Id (MemoryId, SessionId, genMemoryId, genSessionId, idText, parseIdLenient) +import Kioku.Memory qualified as Memory +import Kioku.Memory.Domain + ( ArchiveMemoryData (..), + RecordMemoryData (..), + SupersedeMemoryData (..), + UpdateMemoryConfidenceData (..), + UpdateMemoryTagsData (..), + ) +import Kioku.Memory.Embedding (EmbeddingConfig (..), toEmbeddingModel) +import Kioku.Memory.EventStream (memoryStream) +import Kioku.Migrations.TestSupport (withKiokuMigratedDatabase) +import Kioku.Prelude +import Kioku.Recall.Capability (VectorCapability (..)) +import Kioku.Session qualified as Session +import Kioku.Session.Domain (CompleteSessionData (..), RecordTurnData (..), StartSessionData (..)) +import Kiroku.Store.Connection (defaultConnectionSettings) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Read (readStreamForward) +import Kiroku.Store.Transaction (runTransaction) +import Kiroku.Store.Types (EventType (..), RecordedEvent (..), StreamVersion (..)) +import Shikumi.Effect.Time (runTime) +import Shikumi.Error (ShikumiError (..)) +import Shikumi.LLM (LLM (..)) +import Shikumi.Program (Program, runProgram) +import Shikumi.Schema (Validatable (..)) +import Shikumi.Schema.Types (field, unField) +import Shikumi.Trace (runTrace, tracedLLM) +import Shikumi.Trace.Replay (runLLMReplay) +import Shikumi.Trace.Store (replayIndex) +import System.Directory (doesFileExist) +import System.IO.Temp (withSystemTempDirectory) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Distillation pyramid" + [ testCase "replay distills duplicate turns into merged atom scene and persona" testReplayDistillation, + testCase "re-running distillSessionL1 creates no new memories or audit rows" testRerunIdempotent, + testCase "consolidation failure stores nothing and fails the pass" testConsolidationFailure, + testCase "merge with a missing target drops it and stays convergent" testMergeMissingTarget, + testCase "watermark skips re-extraction until a new turn arrives" testWatermarkSkip, + testCase "a session accumulates one idle timer however many turns" testIdleTimerCollapse, + testCase "recall candidates find a duplicate outside the scan window" testRecallCandidateWindow, + forgetPropagationTests, + confidencePropagationTests, + validationTests + ] + +-- | Forgetting a memory must reach every derived artifact: the scene row, the +-- persona row, and the plaintext mirror files a host agent actually reads. +forgetPropagationTests :: TestTree +forgetPropagationTests = + testGroup + "Forget propagation" + [ testCase "forget operations schedule scene timers" testForgetSchedulesSceneTimers, + testCase "an emptied scope deletes its scene, persona, and mirrors" testEmptyScopeDeletesArtifacts, + testCase "the timer worker propagates an archive to every artifact" testWorkerPropagatesArchive, + testCase "supersede and merge propagate like archive" testWorkerPropagatesSupersedeAndMerge + ] + +-- | A memory's confidence is part of what its scope's scene is built from: it is +-- hashed into the scene's source hash ('atomSource') and written into the LLM's +-- prompt ('renderAtom'). Changing it must therefore refresh the scene, exactly as +-- forgetting does — otherwise an agent that downgrades a belief to @low@ goes on +-- reading a scene that asserts it at @high@. +confidencePropagationTests :: TestTree +confidencePropagationTests = + testGroup + "Confidence propagation" + [ testCase "two confidence changes schedule two distinct scene timers" testConfidenceSchedulesDistinctSceneTimers, + testCase "a tag change schedules nothing, because tags are not in the scene" testTagsUpdateSchedulesNoSceneTimer, + testCase "a confidence change refreshes the scene, the persona, and the mirrors" testWorkerPropagatesConfidence + ] + +-- | The case that guards the trap. A naive fix gives the confidence timer a fixed +-- source id of @\:confidence@, which derives the same UUIDv5 timer id +-- every time — and keiro's scheduling upsert only re-arms a conflicting timer +-- @WHERE status = 'scheduled'@, so the second change would update zero rows and be +-- silently dropped. That fix would refresh the scene on the first change and never +-- again: it would look fixed, and be worse than the bug. The delta-count below is +-- the only thing standing between that mistake and production, so it asserts the +-- count rises on *both* changes, not just the first. +-- +-- The walk is genuinely @high -> medium -> low@ on purpose: 'Memory.updateConfidence' +-- refuses to emit an event when the confidence is unchanged, so re-applying the same +-- value would schedule nothing and pass this test for the wrong reason. +-- +-- No session is written, so every timer counted here is an L2 scene timer. +testConfidenceSchedulesDistinctSceneTimers :: Assertion +testConfidenceSchedulesDistinctSceneTimers = withDistillEnv \env -> do + now <- getCurrentTime + memoryId <- genMemoryId + let scope = forgetScope "intention_confidence_timers" + -- Timers are debounced 5 seconds past their event time; an hour out makes + -- every one of them due. + horizon = addUTCTime 3600 now + result <- + runAppIO env do + recordForgetFixture memoryId scope "Content for the confidence timer test" now + afterRecord <- countDueTimers horizon + + lowered <- + Memory.updateConfidence + UpdateMemoryConfidenceData {memoryId, confidence = MediumConfidence, updatedAt = now} + void (liftIO (expectRight "Memory.updateConfidence to medium" lowered)) + afterFirst <- countDueTimers horizon + + loweredAgain <- + Memory.updateConfidence + UpdateMemoryConfidenceData {memoryId, confidence = LowConfidence, updatedAt = now} + void (liftIO (expectRight "Memory.updateConfidence to low" loweredAgain)) + afterSecond <- countDueTimers horizon + pure (afterRecord, afterFirst, afterSecond) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (afterRecord, afterFirst, afterSecond) -> do + afterRecord @?= 1 + afterFirst @?= 2 + afterSecond @?= 3 + +-- | Tags are in neither 'atomSource' (the scene's source hash) nor 'renderAtom' +-- (the LLM prompt), so a tag change cannot alter the scene. Scheduling a +-- regeneration for one would spend an LLM call to rewrite a byte-identical row. +-- This case pins that reasoning so a future reader does not "complete" the +-- confidence fix by adding an arm for 'MemoryTagsUpdated' too. +testTagsUpdateSchedulesNoSceneTimer :: Assertion +testTagsUpdateSchedulesNoSceneTimer = withDistillEnv \env -> do + now <- getCurrentTime + memoryId <- genMemoryId + let scope = forgetScope "intention_tags_timers" + horizon = addUTCTime 3600 now + result <- + runAppIO env do + recordForgetFixture memoryId scope "Content for the tags timer test" now + afterRecord <- countDueTimers horizon + + -- A genuinely different tag set: 'Memory.updateTags' short-circuits on an + -- unchanged one, which would make this pass for the wrong reason. + retagResult <- + Memory.updateTags + UpdateMemoryTagsData {memoryId, tags = Set.fromList ["retagged"], updatedAt = now} + void (liftIO (expectRight "Memory.updateTags" retagResult)) + afterRetag <- countDueTimers horizon + pure (afterRecord, afterRetag) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (afterRecord, afterRetag) -> do + afterRecord @?= 1 + afterRetag @?= 1 + +-- | The whole pipeline, with nothing called by hand: lowering a memory's confidence +-- schedules a timer through the inline projection, the worker claims and fires it +-- exactly as @kioku worker@ does, and the scene row, the persona row, and the +-- plaintext mirror a host agent actually reads all follow. +testWorkerPropagatesConfidence :: Assertion +testWorkerPropagatesConfidence = withDistillWorkspaceEnv \env workspace -> do + calls <- newDistillCalls + runtime <- echoingRuntime calls <$> replayRuntimeIn workspace + memoryId <- genMemoryId + now <- getCurrentTime + let scope = forgetScope "intention_confidence_worker" + result <- + runAppIO env do + -- 'recordForgetFixture' records at 'HighConfidence'. + recordForgetFixture memoryId scope confidenceContent now + void (drainTimers runtime) + builtScene <- getScenesByScope scope >>= liftIO . expectOneScene "the initial distillation" + builtPersona <- + getPersonaByScope scope >>= liftIO . expectJust "a persona after the initial distillation" + personaRunsBefore <- liftIO (readIORef calls.personaCalls) + + lowered <- + Memory.updateConfidence + UpdateMemoryConfidenceData {memoryId, confidence = LowConfidence, updatedAt = now} + void (liftIO (expectRight "Memory.updateConfidence" lowered)) + void (drainTimers runtime) + + refreshedScene <- + getScenesByScope scope >>= liftIO . expectOneScene "after lowering the confidence" + refreshedPersona <- + getPersonaByScope scope >>= liftIO . expectJust "a persona after lowering the confidence" + personaRunsAfter <- liftIO (readIORef calls.personaCalls) + mirror <- liftIO (TextIO.readFile (sceneMirrorPath workspace refreshedScene)) + + -- Timer firing is at-least-once, so another pass must change nothing. + refired <- drainTimers runtime + pure + ( builtScene, + refreshedScene, + builtPersona, + refreshedPersona, + personaRunsBefore, + personaRunsAfter, + mirror, + refired + ) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (builtScene, refreshedScene, builtPersona, refreshedPersona, personaRunsBefore, personaRunsAfter, mirror, refired) -> do + -- The crisp, LLM-independent proof that the scene was genuinely rebuilt: the + -- source hash is computed over 'atomSource', which contains the confidence. + assertBool + ("the scene kept its stale source hash: " <> Text.unpack refreshedScene.sourceHash) + (refreshedScene.sourceHash /= builtScene.sourceHash) + + -- The prompt the scene distiller actually saw. 'renderAtom' formats each + -- memory as "- (, ): ". + atoms <- latestSceneAtoms calls + assertBool + ("the scene LLM still saw the old confidence: " <> Text.unpack atoms) + (not (highNeedle `Text.isInfixOf` atoms)) + assertBool + ("the scene LLM never saw the new confidence: " <> Text.unpack atoms) + (lowNeedle `Text.isInfixOf` atoms) + + -- ...and, the point of the whole plan, the plaintext file on disk. + assertBool + ("the scene mirror still asserts the old confidence: " <> Text.unpack mirror) + (not (highNeedle `Text.isInfixOf` mirror)) + assertBool + ("the scene mirror never learned the new confidence: " <> Text.unpack mirror) + (lowNeedle `Text.isInfixOf` mirror) + + -- The persona cascades off the scene write. Persona bodies are built from + -- scene titles and bodies, so they inherit the confidence transitively; the + -- observable here is the fact of regeneration, not persona text saying "low". + assertBool + "the persona was not regenerated after the scene changed" + (personaRunsAfter > personaRunsBefore) + assertBool + "the persona row was not rewritten after the scene changed" + (refreshedPersona.updatedAt >= builtPersona.updatedAt) + + refired @?= 0 + +confidenceContent :: Text +confidenceContent = "The user prefers tabs over spaces." + +-- | 'renderAtom' renders "- (, ): ", and +-- 'recordForgetFixture' records a 'MemoryPreference'. Matching the parenthesised +-- pair rather than a bare "high"/"low" keeps the assertion from colliding with +-- any occurrence of those words in the memory's content. +highNeedle :: Text +highNeedle = "(preference, high)" + +lowNeedle :: Text +lowNeedle = "(preference, low)" + +alphaContent :: Text +alphaContent = "The alpha secret is that the launch slipped to March." + +alphaNeedle :: Text +alphaNeedle = "alpha secret" + +betaContent :: Text +betaContent = "The beta fact is that the docs live in ops/README.md." + +betaNeedle :: Text +betaNeedle = "beta fact" + +oldAddressContent :: Text +oldAddressContent = "The user is at 12 Elm Street." + +oldAddressNeedle :: Text +oldAddressNeedle = "12 Elm Street" + +newAddressContent :: Text +newAddressContent = "The user is at 9 Oak Lane." + +newAddressNeedle :: Text +newAddressNeedle = "9 Oak Lane" + +loserContent :: Text +loserContent = "The user mentioned gamma trivia once." + +loserNeedle :: Text +loserNeedle = "gamma trivia" + +winnerContent :: Text +winnerContent = "The user relies on delta truth daily." + +-- | The privacy-shaped case. Forgetting the last memory in a scope must take the +-- scene and persona with it — row and plaintext mirror alike — and must not pay +-- an LLM to summarize what is no longer there. +testEmptyScopeDeletesArtifacts :: Assertion +testEmptyScopeDeletesArtifacts = withDistillWorkspaceEnv \env workspace -> do + calls <- newDistillCalls + runtime <- countingRuntime calls <$> replayRuntimeIn workspace + alphaId <- genMemoryId + betaId <- genMemoryId + now <- getCurrentTime + let scope = forgetScope "intention_forget_empty" + result <- + runAppIO env do + recordForgetFixture alphaId scope alphaContent now + recordForgetFixture betaId scope betaContent now + + firstScene <- regenerateScene runtime scope + sceneRow <- liftIO (expectJustRow "the first regeneration writes a scene" firstScene) + firstPersona <- regeneratePersona runtime scope + personaRow <- liftIO (expectJustRow "the first regeneration writes a persona" firstPersona) + + -- Both mirrors have to be observed here, while they still exist: the whole + -- point of what follows is that they stop existing. + let scenePath = sceneMirrorPath workspace sceneRow + personaPath = personaMirrorPath workspace personaRow + mirrorsWritten <- + liftIO ((&&) <$> doesFileExist scenePath <*> doesFileExist personaPath) + + -- Forget one of the two: the scene must rebuild from the survivor alone. + archivedAlpha <- Memory.archive ArchiveMemoryData {memoryId = alphaId, archivedAt = now} + void (liftIO (expectRight "Memory.archive alpha" archivedAlpha)) + afterOne <- regenerateScene runtime scope + survivorScene <- liftIO (expectJustRow "the scene survives one forget" afterOne) + + -- Forget the last one: nothing may be left to regenerate from. + archivedBeta <- Memory.archive ArchiveMemoryData {memoryId = betaId, archivedAt = now} + void (liftIO (expectRight "Memory.archive beta" archivedBeta)) + emptyScene <- regenerateScene runtime scope + emptyPersona <- regeneratePersona runtime scope + + scenes <- getScenesByScope scope + persona <- getPersonaByScope scope + pure (mirrorsWritten, scenePath, personaPath, survivorScene, emptyScene, emptyPersona, scenes, persona) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (mirrorsWritten, scenePath, personaPath, survivorScene, emptyScene, emptyPersona, scenes, persona) -> do + assertBool "the scene and persona mirrors were written to begin with" mirrorsWritten + + -- The forgotten memory is gone from what the scene is built from. + atoms <- latestSceneAtoms calls + assertBool + ("forgotten content reached the scene LLM: " <> Text.unpack atoms) + (not (alphaNeedle `Text.isInfixOf` atoms)) + survivorScene.atomIds @?= [idText betaId] + + -- The emptied scope keeps no artifact anywhere. + case emptyScene of + Right Nothing -> pure () + other -> assertFailure ("expected no scene for an emptied scope, got " <> show other) + case emptyPersona of + Right Nothing -> pure () + other -> assertFailure ("expected no persona for an emptied scope, got " <> show other) + scenes @?= [] + persona @?= Nothing + doesFileExist scenePath + >>= \exists -> assertBool "the scene mirror was removed" (not exists) + doesFileExist personaPath + >>= \exists -> assertBool "the persona mirror was removed" (not exists) + + -- Neither empty regeneration paid for an LLM call: two scene runs (the + -- initial build and the rebuild from the survivor) and one persona run. + readIORef calls.sceneCalls >>= \runs -> runs @?= 2 + readIORef calls.personaCalls >>= \runs -> runs @?= 1 + +-- | Archive, supersede, and merge each leave behind exactly one scene timer, +-- just as recording does. No session is written, so every timer counted here is +-- an L2 scene timer. +testForgetSchedulesSceneTimers :: Assertion +testForgetSchedulesSceneTimers = withDistillEnv \env -> do + now <- getCurrentTime + archivedId <- genMemoryId + supersededId <- genMemoryId + supersederId <- genMemoryId + loserId <- genMemoryId + winnerId <- genMemoryId + let scope = forgetScope "intention_forget_timers" + -- Timers are debounced 5 seconds past their event time; an hour out makes + -- every one of them due. + horizon = addUTCTime 3600 now + result <- + runAppIO env do + traverse_ + (\mid -> recordForgetFixture mid scope ("Content for " <> idText mid) now) + [archivedId, supersededId, supersederId, loserId, winnerId] + afterRecords <- countDueTimers horizon + + archived <- Memory.archive ArchiveMemoryData {memoryId = archivedId, archivedAt = now} + void (liftIO (expectRight "Memory.archive" archived)) + afterArchive <- countDueTimers horizon + + superseded <- + Memory.supersede + SupersedeMemoryData + { memoryId = supersededId, + supersededBy = supersederId, + supersededAt = now + } + void (liftIO (expectRight "Memory.supersede" superseded)) + afterSupersede <- countDueTimers horizon + + merged <- Memory.merge loserId winnerId + void (liftIO (expectRight "Memory.merge" merged)) + afterMerge <- countDueTimers horizon + pure (afterRecords, afterArchive, afterSupersede, afterMerge) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (afterRecords, afterArchive, afterSupersede, afterMerge) -> do + afterRecords @?= 5 + afterArchive @?= 6 + afterSupersede @?= 7 + afterMerge @?= 8 + +-- | The whole pipeline, with nothing called by hand: a forget command schedules a +-- timer through the inline projection, the worker claims and fires it exactly as +-- @kioku worker@ does, and the scene, persona, and both plaintext mirrors follow. +testWorkerPropagatesArchive :: Assertion +testWorkerPropagatesArchive = withDistillWorkspaceEnv \env workspace -> do + calls <- newDistillCalls + runtime <- echoingRuntime calls <$> replayRuntimeIn workspace + alphaId <- genMemoryId + betaId <- genMemoryId + now <- getCurrentTime + let scope = forgetScope "intention_forget_worker" + result <- + runAppIO env do + recordForgetFixture alphaId scope alphaContent now + recordForgetFixture betaId scope betaContent now + void (drainTimers runtime) + + builtScene <- getScenesByScope scope >>= liftIO . expectOneScene "the initial distillation" + builtPersona <- + getPersonaByScope scope >>= liftIO . expectJust "a persona after the initial distillation" + let scenePath = sceneMirrorPath workspace builtScene + personaPath = personaMirrorPath workspace builtPersona + builtMirrors <- liftIO ((&&) <$> doesFileExist scenePath <*> doesFileExist personaPath) + + -- Forget alpha. The worker has to rebuild the scene from the survivor. + archivedAlpha <- Memory.archive ArchiveMemoryData {memoryId = alphaId, archivedAt = now} + void (liftIO (expectRight "Memory.archive alpha" archivedAlpha)) + void (drainTimers runtime) + survivorScene <- getScenesByScope scope >>= liftIO . expectOneScene "after forgetting alpha" + survivorMirror <- liftIO (TextIO.readFile scenePath) + + -- Forget beta, the last one. Every artifact has to go with it. + archivedBeta <- Memory.archive ArchiveMemoryData {memoryId = betaId, archivedAt = now} + void (liftIO (expectRight "Memory.archive beta" archivedBeta)) + void (drainTimers runtime) + emptyScenes <- getScenesByScope scope + emptyPersona <- getPersonaByScope scope + survivingMirrors <- liftIO ((||) <$> doesFileExist scenePath <*> doesFileExist personaPath) + + -- Timer firing is at-least-once, so another pass must change nothing. + refired <- drainTimers runtime + pure + ( builtMirrors, + survivorScene, + survivorMirror, + emptyScenes, + emptyPersona, + survivingMirrors, + refired + ) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (builtMirrors, survivorScene, survivorMirror, emptyScenes, emptyPersona, survivingMirrors, refired) -> do + assertBool "the worker wrote both mirrors to begin with" builtMirrors + + -- The forgotten memory is gone from the scene's inputs and its metadata... + atoms <- latestSceneAtoms calls + assertBool + ("forgotten content reached the scene LLM: " <> Text.unpack atoms) + (not (alphaNeedle `Text.isInfixOf` atoms)) + survivorScene.atomIds @?= [idText betaId] + + -- ...and, the point of the whole plan, from the plaintext file on disk. + assertBool + ("the forgotten text is still in the scene mirror: " <> Text.unpack survivorMirror) + (not (alphaNeedle `Text.isInfixOf` survivorMirror)) + assertBool + ("the surviving memory vanished from the scene mirror: " <> Text.unpack survivorMirror) + (betaNeedle `Text.isInfixOf` survivorMirror) + + -- Forgetting the last memory empties the scope entirely. + emptyScenes @?= [] + emptyPersona @?= Nothing + assertBool "a mirror file survived an emptied scope" (not survivingMirrors) + + refired @?= 0 + +-- | Archive is not a special case: superseding and merging retire a memory the +-- same way, and must reach the scene the same way. +testWorkerPropagatesSupersedeAndMerge :: Assertion +testWorkerPropagatesSupersedeAndMerge = withDistillWorkspaceEnv \env workspace -> do + supersedeCalls <- newDistillCalls + mergeCalls <- newDistillCalls + supersedeRuntime <- echoingRuntime supersedeCalls <$> replayRuntimeIn workspace + mergeRuntime <- echoingRuntime mergeCalls <$> replayRuntimeIn workspace + oldId <- genMemoryId + newId <- genMemoryId + loserId <- genMemoryId + winnerId <- genMemoryId + now <- getCurrentTime + let supersedeScope = forgetScope "intention_forget_supersede" + mergeScope = forgetScope "intention_forget_merge" + result <- + runAppIO env do + -- Supersession, in its own scope. Drain between the two scopes so each + -- runtime only ever sees its own timers. + recordForgetFixture oldId supersedeScope oldAddressContent now + recordForgetFixture newId supersedeScope newAddressContent now + void (drainTimers supersedeRuntime) + superseded <- + Memory.supersede + SupersedeMemoryData {memoryId = oldId, supersededBy = newId, supersededAt = now} + void (liftIO (expectRight "Memory.supersede" superseded)) + void (drainTimers supersedeRuntime) + supersedeScene <- + getScenesByScope supersedeScope >>= liftIO . expectOneScene "after superseding" + + -- Merge, in a second scope. + recordForgetFixture loserId mergeScope loserContent now + recordForgetFixture winnerId mergeScope winnerContent now + void (drainTimers mergeRuntime) + merged <- Memory.merge loserId winnerId + void (liftIO (expectRight "Memory.merge" merged)) + void (drainTimers mergeRuntime) + mergeScene <- getScenesByScope mergeScope >>= liftIO . expectOneScene "after merging" + + pure (supersedeScene, mergeScene) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (supersedeScene, mergeScene) -> do + supersedeAtoms <- latestSceneAtoms supersedeCalls + assertBool + ("superseded content reached the scene LLM: " <> Text.unpack supersedeAtoms) + (not (oldAddressNeedle `Text.isInfixOf` supersedeAtoms)) + assertBool + ("the superseding content is missing from the scene: " <> Text.unpack supersedeAtoms) + (newAddressNeedle `Text.isInfixOf` supersedeAtoms) + supersedeScene.atomIds @?= [idText newId] + + mergeAtoms <- latestSceneAtoms mergeCalls + assertBool + ("merged-away content reached the scene LLM: " <> Text.unpack mergeAtoms) + (not (loserNeedle `Text.isInfixOf` mergeAtoms)) + mergeScene.atomIds @?= [idText winnerId] + +-- | Fire due timers until none are claimable, looking an hour ahead so both the +-- 5-second scene debounce and the persona timer chained off a scene write are +-- due. Bounded on purpose: a timer that rescheduled itself would otherwise spin +-- here forever, and a hung test is worse than a failed one. +drainTimers :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + DistillRuntime -> + Eff es Int +drainTimers rt = go (50 :: Int) 0 + where + go fuel fired + | fuel <= 0 = liftIO (assertFailure "the timer drain did not converge within 50 fires") + | otherwise = do + realNow <- liftIO getCurrentTime + claimed <- + runKiokuTimerWorkerOnce + Nothing + rt + (scopedScanCandidates 5) + (addUTCTime 3600 realNow) + case claimed of + Nothing -> pure fired + Just _ -> go (fuel - 1) (fired + 1) + +expectOneScene :: String -> [SceneRow] -> IO SceneRow +expectOneScene label = \case + [row] -> pure row + other -> assertFailure (label <> ": expected exactly one scene, got " <> show (length other)) + +expectJust :: String -> Maybe a -> IO a +expectJust label = + maybe (assertFailure (label <> ": expected a value, got none")) pure + +forgetScope :: Text -> MemoryScope +forgetScope = + ScopeEntity (Namespace "rei") (ScopeKind "intention") + +recordForgetFixture :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + MemoryId -> + MemoryScope -> + Text -> + UTCTime -> + Eff es () +recordForgetFixture memoryId scope content now = do + recorded <- + Memory.record + RecordMemoryData + { memoryId, + agentId = "test-agent", + sessionId = Nothing, + scope, + memoryType = MemoryPreference, + content, + priority = 50, + confidence = HighConfidence, + tags = Set.fromList ["forget-test"], + supersedes = Nothing, + recordedAt = now + } + void (liftIO (expectRight "Memory.record" recorded)) + +-- | Pure tests over the 'Validatable' instances; no database, no LLM. +validationTests :: TestTree +validationTests = + testGroup + "extract and consolidate outputs are clamped and rejected" + [ testCase "priority is clamped into [0, 100]" do + validatedAtom (atomWith "preference" "Keep it short." (-1000000) "high") + >>= \atom -> unField atom.priority @?= 0 + validatedAtom (atomWith "preference" "Keep it short." 250 "high") + >>= \atom -> unField atom.priority @?= 100, + testCase "atomType and confidence are lowercased" do + atom <- validatedAtom (atomWith " Preference " "Keep it short." 50 "HIGH") + unField atom.atomType @?= "preference" + unField atom.confidence @?= "high", + testCase "an unknown atomType is rejected" $ + assertValidationFailure "atomType" (ExtractOutput [atomWith "vibe" "Keep it short." 50 "high"]), + testCase "an unknown confidence is rejected" $ + assertValidationFailure "confidence" (ExtractOutput [atomWith "fact" "Keep it short." 50 "certain"]), + testCase "blank atom content is rejected" $ + assertValidationFailure "atom content" (ExtractOutput [atomWith "fact" " " 50 "high"]), + testCase "StoreAtom and SkipAtom drop stray targets" do + stored <- expectValid (decisionWith StoreAtom ["kioku_memory_01hxxx"]) + stored.targetMemoryIds @?= [] + skipped <- expectValid (decisionWith SkipAtom ["garbage"]) + skipped.targetMemoryIds @?= [], + testCase "MergeAtom and UpdateAtom require at least one target" do + assertValidationFailure "MergeAtom requires" (decisionWith MergeAtom []) + assertValidationFailure "UpdateAtom requires" (decisionWith UpdateAtom []), + testCase "an unparseable target id is rejected" $ + assertValidationFailure "unparseable id" (decisionWith MergeAtom ["not-a-typeid"]), + testCase "a parseable target id is accepted unchanged" do + target <- idText <$> genMemoryId + decision <- expectValid (decisionWith MergeAtom [target]) + decision.targetMemoryIds @?= [target] + ] + +atomWith :: Text -> Text -> Int -> Text -> ExtractedAtom +atomWith atomType content priority confidence = + ExtractedAtom + { atomType = field atomType, + content = field content, + priority = field priority, + confidence = field confidence + } + +decisionWith :: ConsolidationAction -> [Text] -> ConsolidationDecision +decisionWith action targetMemoryIds = + ConsolidationDecision + { action, + targetMemoryIds, + resultContent = Nothing, + rationale = field "because" + } + +validatedAtom :: ExtractedAtom -> IO ExtractedAtom +validatedAtom atom = do + output <- expectValid (ExtractOutput [atom]) + case output.atoms of + [validated] -> pure validated + other -> assertFailure ("expected exactly one atom, got " <> show (length other)) + +expectValid :: (Validatable a) => a -> IO a +expectValid value = + case validate value of + Left err -> assertFailure ("expected a valid value, got: " <> Text.unpack err) + Right ok -> pure ok + +assertValidationFailure :: (Validatable a, Show a) => Text -> a -> Assertion +assertValidationFailure needle value = + case validate value of + Right ok -> assertFailure ("expected a validation failure, got: " <> show ok) + Left err -> + assertBool + ("validation message should mention " <> show needle <> ", got: " <> Text.unpack err) + (needle `Text.isInfixOf` err) + +data MemoryStatus = MemoryStatus + { memoryId :: !Text, + content :: !Text, + status :: !Text + } + deriving stock (Generic, Eq, Show) + +data ScopeParams = ScopeParams + { namespace :: !Text, + scopeKind :: !(Maybe Text), + scopeRef :: !(Maybe Text) + } + deriving stock (Generic, Eq, Show) + +data AuditRowView = AuditRowView + { decision :: !Text, + targetIds :: !Text, + resultMemoryId :: !(Maybe Text) + } + deriving stock (Generic, Eq, Show) + +data TimerKindRow = TimerKindRow + { kind :: !Text, + timerCount :: !Int64, + maxFireAt :: !UTCTime + } + deriving stock (Generic, Eq, Show) + +data TimerQuery = TimerQuery + { processManagerName :: !Text, + correlationId :: !Text + } + deriving stock (Generic, Eq, Show) + +data DistillResult = DistillResult + { summary :: !L1Summary, + memories :: ![MemoryStatus], + scenes :: ![SceneRow], + persona :: !(Maybe PersonaRow), + mergeAuditCount :: !Int64, + loserEvents :: ![RecordedEvent] + } + deriving stock (Generic, Show) + +testReplayDistillation :: Assertion +testReplayDistillation = withDistillWorkspaceEnv \env workspace -> do + runtime <- replayRuntimeIn workspace + sid <- genSessionId + now <- getCurrentTime + let scope = fixtureScope + result <- + runAppIO env do + writeFixtureSession sid scope now + distillResult <- distillSessionL1 RespectWatermark runtime (scopedScanCandidates 5) sid + summary <- liftIO (expectDistilled "distillSessionL1" distillResult) + sceneResult <- regenerateScene runtime scope + _scene <- liftIO (expectRight "regenerateScene" sceneResult) + personaResult <- regeneratePersona runtime scope + _persona <- liftIO (expectRight "regeneratePersona" personaResult) + memories <- loadMemoryStatuses scope + mergeAuditCount <- loadMergeAuditCount scope + scenes <- getScenesByScope scope + persona <- getPersonaByScope scope + loserEvents <- loadLoserEvents memories + pure DistillResult {summary, memories, scenes, persona, mergeAuditCount, loserEvents} + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right distill -> assertDistillResult distill + +-- | Run an action against a freshly migrated ephemeral database. Deliberately +-- does not change the working directory: tasty runs test cases concurrently and +-- the process-wide cwd would race. ephemeral-pg keeps its cluster under the XDG +-- cache directory, so no chdir is needed. +withDistillEnv :: (AppEnv -> IO a) -> IO a +withDistillEnv action = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) action + +fixtureScope :: MemoryScope +fixtureScope = + ScopeEntity (Namespace "rei") (ScopeKind "intention") "intention_distill_test" + +-- | A second pass over an unchanged session must converge: the deterministic +-- atom ids mean every write is a no-op and the deterministic audit key means +-- @ON CONFLICT DO NOTHING@ suppresses duplicate audit rows. +testRerunIdempotent :: Assertion +testRerunIdempotent = withDistillEnv \env -> do + runtime <- replayRuntime + sid <- genSessionId + now <- getCurrentTime + result <- + runAppIO env do + writeFixtureSession sid fixtureScope now + first <- distillSessionL1 RespectWatermark runtime (scopedScanCandidates 5) sid + summary1 <- liftIO (expectDistilled "first pass" first) + memories1 <- loadMemoryStatuses fixtureScope + audits1 <- loadAuditCount fixtureScope + second <- distillSessionL1 IgnoreWatermark runtime (scopedScanCandidates 5) sid + summary2 <- liftIO (expectDistilled "second pass" second) + memories2 <- loadMemoryStatuses fixtureScope + audits2 <- loadAuditCount fixtureScope + pure (summary1, memories1, audits1, summary2, memories2, audits2) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (summary1, memories1, audits1, summary2, memories2, audits2) -> do + summary1.stored @?= 1 + summary1.merged @?= 1 + length memories1 @?= 2 + audits1 @?= 2 + summary2.extracted @?= 2 + summary2.stored @?= 0 + summary2.merged @?= 0 + summary2.skipped @?= 2 + memories2 @?= memories1 + audits2 @?= audits1 + length [() | row <- memories2, row.status == "active"] @?= 1 + +-- | A consolidation LLM failure is a failed pass, not a silent store. +testConsolidationFailure :: Assertion +testConsolidationFailure = withDistillEnv \env -> do + base <- replayRuntime + let runtime = base {runConsolidate = \_ -> pure (Left (ValidationFailure "boom"))} + sid <- genSessionId + now <- getCurrentTime + result <- + runAppIO env do + writeFixtureSession sid fixtureScope now + distilled <- distillSessionL1 RespectWatermark runtime (scopedScanCandidates 5) sid + memories <- loadMemoryStatuses fixtureScope + audits <- loadAuditCount fixtureScope + pure (distilled, memories, audits) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (distilled, memories, audits) -> do + case distilled of + Left (L1ConsolidationFailed _) -> pure () + other -> assertFailure ("expected L1ConsolidationFailed, got " <> show other) + memories @?= [] + audits @?= 0 + +-- | A hallucinated-but-parseable merge target is dropped before the winner is +-- recorded, the audit records only the surviving target, and a re-run converges. +testMergeMissingTarget :: Assertion +testMergeMissingTarget = withDistillEnv \env -> do + base <- replayRuntime + existingId <- genMemoryId + ghostId <- genMemoryId + let mergeResponse = mergeTargetsResponse [idText existingId, idText ghostId] + runtime = + base + { runExtract = replayProgram singleAtomExtractResponse extractProgram, + runConsolidate = replayProgram mergeResponse consolidateProgram + } + sid <- genSessionId + now <- getCurrentTime + result <- + runAppIO env do + writeFixtureSession sid fixtureScope now + seedMemory existingId sid fixtureScope "The user likes concise answers." now + first <- distillSessionL1 RespectWatermark runtime (scopedScanCandidates 5) sid + summary1 <- liftIO (expectDistilled "first pass" first) + memories1 <- loadMemoryStatuses fixtureScope + audits1 <- loadAuditRows fixtureScope + second <- distillSessionL1 IgnoreWatermark runtime (scopedScanCandidates 5) sid + void (liftIO (expectDistilled "second pass" second)) + memories2 <- loadMemoryStatuses fixtureScope + audits2 <- loadAuditRows fixtureScope + pure (summary1, memories1, audits1, memories2, audits2) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (summary1, memories1, audits1, memories2, audits2) -> do + summary1.extracted @?= 1 + summary1.merged @?= 1 + assertBool "seeded memory is merged away" $ + any (\row -> row.memoryId == idText existingId && row.status == "merged") memories1 + length [() | row <- memories1, row.status == "active"] @?= 1 + case audits1 of + [audit] -> do + audit.decision @?= "merge" + audit.targetIds @?= encodeJsonText [idText existingId] + other -> assertFailure ("expected exactly one audit row, got " <> show (length other)) + memories2 @?= memories1 + audits2 @?= audits1 + +-- | The watermark makes a re-fire with no new turns a pure database read. The +-- proof is an extractor that fails if it is ever called: the pass succeeds +-- anyway, so the LLM was not reached. Recording one more turn re-enables it. +testWatermarkSkip :: Assertion +testWatermarkSkip = withDistillEnv \env -> do + working <- replayRuntime + let exploding = + working {runExtract = \_ -> pure (Left (ValidationFailure "extractor must not run"))} + sid <- genSessionId + now <- getCurrentTime + result <- + runAppIO env do + writeRunningFixtureSession sid fixtureScope now + first <- distillSessionL1 RespectWatermark working (scopedScanCandidates 5) sid + void (liftIO (expectDistilled "first pass" first)) + skipped <- distillSessionL1 RespectWatermark exploding (scopedScanCandidates 5) sid + recordFixtureTurn sid now 4 "Never deploy on a Friday." + afterNewTurn <- distillSessionL1 RespectWatermark exploding (scopedScanCandidates 5) sid + pure (skipped, afterNewTurn) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (skipped, afterNewTurn) -> do + case skipped of + Right L1SkippedUpToDate -> pure () + other -> assertFailure ("expected L1SkippedUpToDate, got " <> show other) + case afterNewTurn of + Left (L1ExtractionFailed _) -> pure () + other -> assertFailure ("expected L1ExtractionFailed after a new turn, got " <> show other) + +-- | However many turns a session records, it holds exactly one idle timer, +-- re-armed forward to the latest turn. Ramp timers fire only on ramp turns +-- (1, 2, 4, 8, 16, ...), and completion adds one final timer. +testIdleTimerCollapse :: Assertion +testIdleTimerCollapse = withDistillEnv \env -> do + sid <- genSessionId + now <- getCurrentTime + result <- + runAppIO env do + writeFixtureSession sid fixtureScope now + loadTimerKinds sid + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right kinds -> do + let lookupKind k = [row | row <- kinds, row.kind == k] + case lookupKind "idle" of + [idle] -> do + idle.timerCount @?= 1 + let lastTurnAt = turnRecordedAt now (length fixtureTurns) + expected = addUTCTime idleFlushSeconds lastTurnAt + assertBool + ( "idle timer should be re-armed to the last turn's recordedAt + 30min; expected " + <> show expected + <> " but found " + <> show idle.maxFireAt + ) + (abs (diffUTCTime idle.maxFireAt expected) < 0.001) + other -> assertFailure ("expected exactly one idle timer row, got " <> show other) + case lookupKind "ramp" of + [ramp] -> ramp.timerCount @?= 2 + other -> assertFailure ("expected one ramp kind grouping, got " <> show other) + case lookupKind "final" of + [final] -> final.timerCount @?= 1 + other -> assertFailure ("expected one final timer row, got " <> show other) + +-- | The scan finder returns the first @limit@ rows of a @priority ASC@ scan, so +-- a duplicate that sits behind enough low-priority filler is invisible to the +-- consolidator and gets stored again. Relevance-ranked recall finds it. +-- +-- Both halves run in one database against separate scopes, so the contrast is +-- the only difference between them. +testRecallCandidateWindow :: Assertion +testRecallCandidateWindow = withDistillEnv \env -> do + base <- replayRuntime + recallSid <- genSessionId + scanSid <- genSessionId + recallDuplicateId <- genMemoryId + scanDuplicateId <- genMemoryId + now <- getCurrentTime + let recallScope = ScopeEntity (Namespace "rei") (ScopeKind "intention") "intention_recall_window" + scanScope = ScopeEntity (Namespace "rei") (ScopeKind "intention") "intention_scan_window" + -- Merge only when the consolidator was actually shown the duplicate. + runtimeFor duplicateId = + base + { runExtract = replayProgram singleAtomExtractResponse extractProgram, + runConsolidate = \input -> + if any (\existing -> unField existing.memoryId == idText duplicateId) input.existing + then replayProgram (mergeTargetsResponse [idText duplicateId]) consolidateProgram input + else replayProgram storeAtomResponse consolidateProgram input + } + -- Inject the capability rather than probing the cluster. This case is about the candidate + -- finder -- that recall reaches a duplicate the priority scan window hides -- and nothing + -- about vectors. Pinning it to the keyword plan is what makes dummyEmbeddingModel safe: it + -- guarantees no embedding endpoint is ever called, instead of relying on the test cluster + -- happening to lack pgvector, which it no longer does. + let capability = VectorExtensionUnavailable + result <- + runAppIO env do + writeRunningFixtureSession recallSid recallScope now + seedCandidateWindow recallSid recallScope now recallDuplicateId + recallOutcome <- + distillSessionL1 + RespectWatermark + (runtimeFor recallDuplicateId) + (recallCandidates dummyEmbeddingModel capability 8) + recallSid + recallSummary <- liftIO (expectDistilled "recall pass" recallOutcome) + recallMemories <- loadMemoryStatuses recallScope + + writeRunningFixtureSession scanSid scanScope now + seedCandidateWindow scanSid scanScope now scanDuplicateId + scanOutcome <- + distillSessionL1 + RespectWatermark + (runtimeFor scanDuplicateId) + (scopedScanCandidates 5) + scanSid + scanSummary <- liftIO (expectDistilled "scan pass" scanOutcome) + scanMemories <- loadMemoryStatuses scanScope + pure (recallSummary, recallMemories, scanSummary, scanMemories) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (recallSummary, recallMemories, scanSummary, scanMemories) -> do + recallSummary.merged @?= 1 + recallSummary.stored @?= 0 + assertBool "recall found the duplicate behind the scan window and merged it" $ + any + (\row -> row.memoryId == idText recallDuplicateId && row.status == "merged") + recallMemories + + -- The defect being fixed, pinned in place: the scan never sees the duplicate. + scanSummary.stored @?= 1 + scanSummary.merged @?= 0 + assertBool "the scan window hid the duplicate, so it stayed active" $ + any + (\row -> row.memoryId == idText scanDuplicateId && row.status == "active") + scanMemories + +-- | Six low-priority fillers ahead of the duplicate in a @priority ASC@ scan, +-- with content that shares no stem with the extracted atom so full-text search +-- passes over them. +seedCandidateWindow :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + MemoryScope -> + UTCTime -> + MemoryId -> + Eff es () +seedCandidateWindow sid scope now duplicateId = do + traverse_ + ( \(n :: Int) -> do + fillerId <- genMemoryId + seedMemoryWith + fillerId + sid + scope + ("Filler note " <> Text.pack (show n) <> ": the deployment pipeline runs on Nix flakes.") + 10 + now + ) + [1 .. 6] + seedMemoryWith duplicateId sid scope "The user prefers concise answers." 90 now + +dummyEmbeddingModel :: EmbeddingModel +dummyEmbeddingModel = + toEmbeddingModel + EmbeddingConfig + { baseUrl = "http://embedding.invalid", + model = "test-embedding", + dimensions = 8, + apiKey = "" + } + +seedMemory :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + MemoryId -> + SessionId -> + MemoryScope -> + Text -> + UTCTime -> + Eff es () +seedMemory memoryId sid scope content = + seedMemoryWith memoryId sid scope content 50 + +seedMemoryWith :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + MemoryId -> + SessionId -> + MemoryScope -> + Text -> + Int -> + UTCTime -> + Eff es () +seedMemoryWith memoryId sid scope content priority now = do + recorded <- + Memory.record + RecordMemoryData + { memoryId, + agentId = "test-agent", + sessionId = Just sid, + scope, + memoryType = MemoryPreference, + content, + priority, + confidence = HighConfidence, + tags = Set.fromList ["seed"], + supersedes = Nothing, + recordedAt = now + } + void (liftIO (expectRight "Memory.record" recorded)) + +-- | Turn @n@ is recorded one minute after turn @n-1@, so the single debounced +-- idle timer's @fire_at@ is observably re-armed forward by each turn. +turnRecordedAt :: UTCTime -> Int -> UTCTime +turnRecordedAt now turnIndex = + addUTCTime (60 * fromIntegral (turnIndex - 1)) now + +fixtureTurns :: [(Int, Text)] +fixtureTurns = + [ (1, "I like short answers."), + (2, "Please keep replies brief."), + (3, "The deploy script is in ops/deploy.sh.") + ] + +-- | Start a session and record 'fixtureTurns', leaving it @Running@ so further +-- turns can be recorded. The aggregate only accepts @RecordTurn@ from @Running@. +writeRunningFixtureSession :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + MemoryScope -> + UTCTime -> + Eff es () +writeRunningFixtureSession sid scope now = do + startResult <- + Session.start + StartSessionData + { sessionId = sid, + agentId = "test-agent", + focus = "style preference capture", + scope, + subjectRef = Just "intention_distill_test", + previousSessionId = Nothing, + parentSessionId = Nothing, + delegationDepth = 0, + startedAt = now + } + void (liftIO (expectRight "Session.start" startResult)) + traverse_ (uncurry (recordFixtureTurn sid now)) fixtureTurns + +recordFixtureTurn :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + UTCTime -> + Int -> + Text -> + Eff es () +recordFixtureTurn sid now turnIndex content = do + turnResult <- + Session.recordTurn + RecordTurnData + { sessionId = sid, + turnId = idText sid <> "-turn-" <> Text.pack (show turnIndex), + turnIndex, + role = "user", + content, + toolSummary = Nothing, + promptTokens = Nothing, + outputTokens = Nothing, + recordedAt = turnRecordedAt now turnIndex + } + void (liftIO (expectRight "Session.recordTurn" turnResult)) + +writeFixtureSession :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + MemoryScope -> + UTCTime -> + Eff es () +writeFixtureSession sid scope now = do + writeRunningFixtureSession sid scope now + completeResult <- + Session.complete + CompleteSessionData + { sessionId = sid, + completedAt = turnRecordedAt now (length fixtureTurns), + modelUsed = Just "test-model", + summary = Just "Captured style preferences." + } + void (liftIO (expectRight "Session.complete" completeResult)) + +replayRuntime :: IO DistillRuntime +replayRuntime = do + rt <- newDistillRuntime + pure + rt + { runExtract = replayProgram extractResponse extractProgram, + runConsolidate = \input -> replayProgram (consolidateResponse input) consolidateProgram input, + runScene = replayProgram sceneResponse sceneProgram, + runPersona = replayProgram personaResponse personaProgram + } + +-- | A replay runtime whose plaintext mirrors go to a private directory rather +-- than the process's working directory. Any test that asserts on mirror files — +-- or merely regenerates a scene, which writes one as a side effect — must use +-- this: tasty runs cases concurrently, so a process-wide @chdir@ would race. +replayRuntimeIn :: FilePath -> IO DistillRuntime +replayRuntimeIn workspace = do + rt <- replayRuntime + pure rt {workspaceRoot = Just workspace} + +-- | An 'AppEnv' plus the private workspace its mirrors are written into. +withDistillWorkspaceEnv :: (AppEnv -> FilePath -> IO a) -> IO a +withDistillWorkspaceEnv action = + withSystemTempDirectory "kioku-distill" \workspace -> + withDistillEnv \env -> action env workspace + +-- | Counters and captures over the two distillation runners the forget paths +-- must not invoke, plus the atom text each scene distillation was actually shown. +data DistillCalls = DistillCalls + { sceneCalls :: !(IORef Int), + personaCalls :: !(IORef Int), + sceneAtoms :: !(IORef [Text]) + } + +newDistillCalls :: IO DistillCalls +newDistillCalls = + DistillCalls <$> newIORef 0 <*> newIORef 0 <*> newIORef [] + +countingRuntime :: DistillCalls -> DistillRuntime -> DistillRuntime +countingRuntime calls rt = + rt + { runScene = \input -> do + modifyIORef' calls.sceneCalls (+ 1) + modifyIORef' calls.sceneAtoms (<> [unField input.atoms]) + rt.runScene input, + runPersona = \input -> do + modifyIORef' calls.personaCalls (+ 1) + rt.runPersona input + } + +-- | 'countingRuntime', but the scene body echoes the atoms it was built from, so +-- the mirror file's bytes on disk are a direct function of which memories +-- survive. That is what lets the end-to-end tests assert "the forgotten text is +-- gone from the file a host agent would read" against real content, rather than +-- settling for the row metadata. +echoingRuntime :: DistillCalls -> DistillRuntime -> DistillRuntime +echoingRuntime calls rt = + (countingRuntime calls rt) + { runScene = \input -> do + modifyIORef' calls.sceneCalls (+ 1) + modifyIORef' calls.sceneAtoms (<> [unField input.atoms]) + replayProgram (echoSceneResponse (unField input.atoms)) sceneProgram input + } + +-- | Newlines are flattened because the response format is line-oriented: the +-- atoms are a bulleted list, and each bullet would otherwise look like a field. +echoSceneResponse :: Text -> Text +echoSceneResponse atoms = + Text.unlines + [ "[[ ## title ## ]]", + "Response style", + "[[ ## bodyMd ## ]]", + Text.replace "\n" " | " atoms, + "[[ ## completed ## ]]" + ] + +-- | The atom text handed to the most recent scene distillation. +latestSceneAtoms :: DistillCalls -> IO Text +latestSceneAtoms calls = do + captured <- readIORef calls.sceneAtoms + case reverse captured of + latest : _ -> pure latest + [] -> assertFailure "the scene distiller was never called" + +replayProgram :: Text -> Program i o -> i -> IO (Either ShikumiError o) +replayProgram response prog input = do + (live, tree) <- + runEff + . runPrim + . runTime + . runTrace + . runFixedLLM (mkResponse response) + . tracedLLM + . runErrorNoCallStack @ShikumiError + $ runProgram prog input + case live of + Left err -> pure (Left err) + Right _ -> case replayIndex tree of + Left err -> assertFailure (Text.unpack err) + Right idx -> + runEff + . runLLMReplay idx + . runErrorNoCallStack @ShikumiError + $ runProgram prog input + +runFixedLLM :: Response -> Eff (LLM : es) a -> Eff es a +runFixedLLM resp = interpret \_ -> \case + Complete {} -> pure resp + Stream {} -> pure [] + +mkResponse :: Text -> Response +mkResponse responseText = + _Response + & #message + . #content + .~ Vector.singleton (AssistantText (_TextContent & #text .~ responseText)) + +extractResponse :: Text +extractResponse = + """ + [[ ## atoms ## ]] + [ + {"atomType":"preference","content":"The user prefers concise answers.","priority":50,"confidence":"high"}, + {"atomType":"preference","content":"The user wants replies kept brief.","priority":50,"confidence":"high"} + ] + [[ ## completed ## ]] + """ + +storeAtomResponse :: Text +storeAtomResponse = + """ + [[ ## action ## ]] + StoreAtom + [[ ## targetMemoryIds ## ]] + [] + [[ ## resultContent ## ]] + The user prefers concise answers. + [[ ## rationale ## ]] + The preference is durable and not yet represented. + [[ ## completed ## ]] + """ + +consolidateResponse :: ConsolidateInput -> Text +consolidateResponse input = + case input.existing of + [] -> storeAtomResponse + ExistingMemory {memoryId = targetId} : _ -> + Text.unlines + [ "[[ ## action ## ]]", + "MergeAtom", + "[[ ## targetMemoryIds ## ]]", + encodeJsonText ([unField targetId] :: [Text]), + "[[ ## resultContent ## ]]", + "The user prefers concise answers.", + "[[ ## rationale ## ]]", + "The candidate restates the existing concise-answer preference.", + "[[ ## completed ## ]]" + ] + +singleAtomExtractResponse :: Text +singleAtomExtractResponse = + """ + [[ ## atoms ## ]] + [ + {"atomType":"preference","content":"The user prefers concise answers.","priority":50,"confidence":"high"} + ] + [[ ## completed ## ]] + """ + +mergeTargetsResponse :: [Text] -> Text +mergeTargetsResponse targets = + Text.unlines + [ "[[ ## action ## ]]", + "MergeAtom", + "[[ ## targetMemoryIds ## ]]", + encodeJsonText targets, + "[[ ## resultContent ## ]]", + "The user prefers concise answers.", + "[[ ## rationale ## ]]", + "The candidate restates the existing concise-answer preference.", + "[[ ## completed ## ]]" + ] + +sceneResponse :: Text +sceneResponse = + """ + [[ ## title ## ]] + Response style + [[ ## bodyMd ## ]] + - The user prefers concise answers. + [[ ## completed ## ]] + """ + +personaResponse :: Text +personaResponse = + """ + [[ ## bodyMd ## ]] + # Persona + + The user values concise answers in this scope. + [[ ## completed ## ]] + """ + +loadMemoryStatuses :: + (Store :> es) => + MemoryScope -> + Eff es [MemoryStatus] +loadMemoryStatuses scope = + runTransaction $ + Tx.statement (scopeParams scope) selectMemoryStatusesStmt + +loadMergeAuditCount :: + (Store :> es) => + MemoryScope -> + Eff es Int64 +loadMergeAuditCount scope = + runTransaction $ + Tx.statement (scopeParams scope) selectMergeAuditCountStmt + +loadAuditCount :: + (Store :> es) => + MemoryScope -> + Eff es Int64 +loadAuditCount scope = + runTransaction $ + Tx.statement (scopeParams scope) selectAuditCountStmt + +loadAuditRows :: + (Store :> es) => + MemoryScope -> + Eff es [AuditRowView] +loadAuditRows scope = + runTransaction $ + Tx.statement (scopeParams scope) selectAuditRowsStmt + +-- | keiro's timer table is unqualified at the pinned version: it lives in the +-- @kiroku@ schema, which the connection's search_path already resolves. +loadTimerKinds :: + (Store :> es) => + SessionId -> + Eff es [TimerKindRow] +loadTimerKinds sid = + runTransaction $ + Tx.statement + TimerQuery + { processManagerName = l1ExtractProcessManagerName, + correlationId = idText sid + } + selectTimerKindsStmt + +selectTimerKindsStmt :: Statement TimerQuery [TimerKindRow] +selectTimerKindsStmt = + preparable + """ + SELECT payload->>'kind', count(*), max(fire_at) + FROM keiro.keiro_timers + WHERE process_manager_name = $1 + AND correlation_id = $2 + GROUP BY payload->>'kind' + ORDER BY payload->>'kind' + """ + timerQueryEncoder + (D.rowList timerKindRowDecoder) + +timerQueryEncoder :: E.Params TimerQuery +timerQueryEncoder = + ((\q -> q.processManagerName) >$< E.param (E.nonNullable E.text)) + <> ((\q -> q.correlationId) >$< E.param (E.nonNullable E.text)) + +timerKindRowDecoder :: D.Row TimerKindRow +timerKindRowDecoder = + TimerKindRow + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.int8) + <*> D.column (D.nonNullable D.timestamptz) + +loadLoserEvents :: + (Store :> es) => + [MemoryStatus] -> + Eff es [RecordedEvent] +loadLoserEvents memories = + case [mid | MemoryStatus {memoryId = mid, status = "merged"} <- memories] of + loser : _ -> + case (parseIdLenient loser :: Either Text MemoryId) of + Left _ -> pure [] + Right loserId -> + Vector.toList <$> readStreamForward (Stream.streamName (memoryStream loserId)) (StreamVersion 0) 100 + [] -> pure [] + +scopeParams :: MemoryScope -> ScopeParams +scopeParams scope = + ScopeParams + { namespace = scopeNamespaceText scope, + scopeKind = scopeKindText scope, + scopeRef = scopeRefText scope + } + +selectMemoryStatusesStmt :: Statement ScopeParams [MemoryStatus] +selectMemoryStatusesStmt = + preparable + """ + SELECT memory_id, content, status + FROM kioku_memories + WHERE namespace = $1 + AND ((scope_kind = $2 AND scope_ref = $3) + OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) + ORDER BY created_at ASC + """ + scopeParamsEncoder + (D.rowList memoryStatusDecoder) + +selectAuditCountStmt :: Statement ScopeParams Int64 +selectAuditCountStmt = + preparable + """ + SELECT count(*) + FROM kioku_consolidation_decisions + WHERE namespace = $1 + AND ((scope_kind = $2 AND scope_ref = $3) + OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) + """ + scopeParamsEncoder + (D.singleRow (D.column (D.nonNullable D.int8))) + +selectAuditRowsStmt :: Statement ScopeParams [AuditRowView] +selectAuditRowsStmt = + preparable + """ + SELECT decision, target_ids::text, result_memory_id + FROM kioku_consolidation_decisions + WHERE namespace = $1 + AND ((scope_kind = $2 AND scope_ref = $3) + OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) + ORDER BY decided_at ASC + """ + scopeParamsEncoder + (D.rowList auditRowViewDecoder) + +auditRowViewDecoder :: D.Row AuditRowView +auditRowViewDecoder = + AuditRowView + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nullable D.text) + +selectMergeAuditCountStmt :: Statement ScopeParams Int64 +selectMergeAuditCountStmt = + preparable + """ + SELECT count(*) + FROM kioku_consolidation_decisions + WHERE namespace = $1 + AND ((scope_kind = $2 AND scope_ref = $3) + OR ($2 IS NULL AND scope_kind IS NULL AND $3 IS NULL AND scope_ref IS NULL)) + AND decision = 'merge' + """ + scopeParamsEncoder + (D.singleRow (D.column (D.nonNullable D.int8))) + +scopeParamsEncoder :: E.Params ScopeParams +scopeParamsEncoder = + ((\q -> q.namespace) >$< E.param (E.nonNullable E.text)) + <> ((\q -> q.scopeKind) >$< E.param (E.nullable E.text)) + <> ((\q -> q.scopeRef) >$< E.param (E.nullable E.text)) + +memoryStatusDecoder :: D.Row MemoryStatus +memoryStatusDecoder = + MemoryStatus + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.text) + +assertDistillResult :: DistillResult -> Assertion +assertDistillResult result = do + result.summary.extracted @?= 2 + result.summary.stored @?= 1 + result.summary.merged @?= 1 + length [() | row <- result.memories, row.status == "active"] @?= 1 + length [() | row <- result.memories, row.status == "merged"] @?= 1 + assertBool "active memory keeps the concise-answer preference" $ + any (\row -> row.status == "active" && "concise answers" `Text.isInfixOf` row.content) result.memories + result.mergeAuditCount @?= 1 + assertBool "loser stream contains MemoryMerged" $ + any (\event -> event.eventType == EventType "MemoryMerged") result.loserEvents + length result.scenes @?= 1 + assertBool "scene body is non-empty" $ + any (not . Text.null . (.bodyMd)) result.scenes + case result.persona of + Nothing -> assertFailure "expected persona row" + Just persona -> assertBool "persona body is non-empty" (not (Text.null persona.bodyMd)) + +expectRight :: (Show e) => String -> Either e a -> IO a +expectRight label = \case + Left err -> assertFailure (label <> " failed: " <> show err) + Right value -> pure value + +-- | Unwrap a regeneration that was expected to produce an artifact, not to find +-- the scope empty. +expectJustRow :: (Show e) => String -> Either e (Maybe a) -> IO a +expectJustRow label outcome = + expectRight label outcome >>= \case + Just row -> pure row + Nothing -> assertFailure (label <> ": expected a row, got none") + +-- | Unwrap a pass that was expected to actually run, not skip on the watermark. +expectDistilled :: String -> Either L1Error L1Outcome -> IO L1Summary +expectDistilled label outcome = do + distilled <- expectRight label outcome + case distilled of + L1Distilled summary -> pure summary + L1SkippedUpToDate -> assertFailure (label <> " unexpectedly skipped on the watermark") + +encodeJsonText :: (ToJSON a) => a -> Text +encodeJsonText = + TE.decodeUtf8 . BL.toStrict . Aeson.encode diff --git a/test/Kioku/EmbeddingWorkerSpec.hs b/test/Kioku/EmbeddingWorkerSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/EmbeddingWorkerSpec.hs @@ -0,0 +1,298 @@ +{-# LANGUAGE DataKinds #-} + +module Kioku.EmbeddingWorkerSpec + ( tests, + ) +where + +import Baikai.Embedding (EmbeddingModel) +import Data.Aeson qualified as Aeson +import Data.HashMap.Strict qualified as HashMap +import Data.Set qualified as Set +import Data.Vector qualified as Vector +import Effectful (Eff, IOE, (:>)) +import Effectful.Error.Static (Error) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Keiro.Stream qualified as Stream +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), ScopeKind (..)) +import Kioku.Api.Types (Confidence (..), MemoryType (..)) +import Kioku.App (AppEffects, AppEnv, runAppIO, withNoopAppEnv) +import Kioku.Id (MemoryId, genMemoryId, idText) +import Kioku.Memory qualified as Memory +import Kioku.Memory.Domain (RecordMemoryData (..)) +import Kioku.Memory.Embedding (EmbedError (..), EmbeddingConfig (..), toEmbeddingModel) +import Kioku.Memory.Embedding.Worker + ( EmbeddingWorkerEnv (..), + embeddingHandler, + shouldSkipEmbedding, + ) +import Kioku.Memory.EventStream (memoryStream) +import Kioku.Migrations.TestSupport (withKiokuMigratedDatabase) +import Kioku.Prelude +import Kioku.Recall.Capability (VectorCapability (..), detectVectorCapability) +import Kioku.Worker.Failure (embeddingRetryDelay, isTransientStoreError) +import Kiroku.Store.Connection (defaultConnectionSettings) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError (..)) +import Kiroku.Store.Read (readStreamForward) +import Kiroku.Store.Transaction (runTransaction) +import Kiroku.Store.Types (ExpectedVersion (..), RecordedEvent (..), StreamName (..), StreamVersion (..)) +import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..), HaltReason (..), RetryDelay (..)) +import Shibuya.Core.AckHandle (AckHandle (..)) +import Shibuya.Core.Ingested (Ingested (..)) +import Shibuya.Core.Types (Attempt (..), Envelope (..), MessageId (..)) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Embedding worker" + [ testCase "skips only when the embedding exists and the content hash matches" do + shouldSkipEmbedding True (Just "hash-a") "hash-a" @?= True + shouldSkipEmbedding False (Just "hash-a") "hash-a" @?= False + shouldSkipEmbedding True Nothing "hash-a" @?= False + shouldSkipEmbedding True (Just "hash-b") "hash-a" @?= False, + testCase "classifies only connection-shaped store errors as transient" testTransientClassification, + testCase "backs off on the documented retry schedule" testRetrySchedule, + testCase "provider failure acks retry" testProviderFailureRetries, + testCase "undecodable payload acks dead-letter" testUndecodablePayloadDeadLetters, + testCase "successful embedding acks ok and stores the vector" testSuccessStoresEmbedding, + testCase "dimension mismatch halts the processor" testDimensionMismatchHalts + ] + +-- | Every constructor kiroku documents as retryable is transient; every +-- constructor describing a stable disagreement with the database is not. A new +-- kiroku constructor breaks this test's exhaustiveness at the source, not here. +testTransientClassification :: Assertion +testTransientClassification = do + isTransientStoreError PoolAcquisitionTimeout @?= True + isTransientStoreError (ConnectionLost "reset by peer") @?= True + isTransientStoreError (ConnectionError "pool closed") @?= True + isTransientStoreError (UnexpectedServerError "22000" "expected 1536 dimensions, not 8") @?= False + isTransientStoreError (StreamNotFound (StreamName "kioku_memory-x")) @?= False + isTransientStoreError (DuplicateEvent Nothing) @?= False + isTransientStoreError (WrongExpectedVersion (StreamName "s") AnyVersion (StreamVersion 0)) @?= False + isTransientStoreError (EmptyAppendBatch (StreamName "s")) @?= False + isTransientStoreError (ReservedStreamName (StreamName "$all")) @?= False + isTransientStoreError (StreamNameTooLong (StreamName "s") 9000) @?= False + isTransientStoreError (StreamAlreadyExists (StreamName "s")) @?= False + isTransientStoreError (EventAlreadyLinked (StreamName "s") Nothing) @?= False + isTransientStoreError (LinkSourceEventMissing (StreamName "s")) @?= False + +testRetrySchedule :: Assertion +testRetrySchedule = do + embeddingRetryDelay (Just (Attempt 0)) @?= RetryDelay 5 + embeddingRetryDelay (Just (Attempt 1)) @?= RetryDelay 20 + embeddingRetryDelay (Just (Attempt 2)) @?= RetryDelay 60 + embeddingRetryDelay (Just (Attempt 3)) @?= RetryDelay 180 + -- An adapter that does not track redeliveries gets the longest delay. + embeddingRetryDelay Nothing @?= RetryDelay 180 + +-- | A provider outage must not be acked as success: the event is redelivered +-- with the attempt-indexed backoff, and kiroku's retry policy bounds it. +testProviderFailureRetries :: Assertion +testProviderFailureRetries = + withVectorEnv "provider failure" \env capability -> do + decision <- + runHandler env capability (failingEmbed (EmbedTransport "provider is down")) (Just 0) + decision @?= AckRetry (embeddingRetryDelay (Just (Attempt 0))) + +-- | A payload that cannot be decoded can never succeed. It goes to the +-- dead-letter table with a reason instead of vanishing behind an 'AckOk'. +-- Not gated on pgvector: the handler decides before it touches the memory row. +testUndecodablePayloadDeadLetters :: Assertion +testUndecodablePayloadDeadLetters = + withEmbeddingEnv \appEnv -> do + capability <- runOrFail appEnv (detectVectorCapability embeddingDims) + decision <- runOrFail appEnv do + (_, recorded) <- recordFixtureMemory "corrupt payload memory" + let env = mkTestEnv (failingEmbed EmbedEmpty) + embeddingHandler capability env (mkIngested (corruptPayload recorded) (Just 0)) + case decision of + AckDeadLetter (InvalidPayload _) -> pure () + other -> assertFailure ("expected a dead-letter for an undecodable payload, got: " <> show other) + +-- | The happy path still works, and it actually writes: 'AckOk' is only honest +-- if the vector landed in the row. +testSuccessStoresEmbedding :: Assertion +testSuccessStoresEmbedding = + withVectorEnv "successful embedding" \env capability -> do + (decision, stored) <- + runOrFail env do + (memoryId, recorded) <- recordFixtureMemory "a memory worth embedding" + let workerEnv = mkTestEnv (\_ -> pure (Right (Vector.replicate embeddingDims 0.1))) + decision <- embeddingHandler capability workerEnv (mkIngested recorded (Just 0)) + stored <- loadEmbeddingState (idText memoryId) + pure (decision, stored) + decision @?= AckOk + assertBool "the memory row has an embedding and a content hash" stored + +-- | The one case where halting is right. A dimension mismatch is a permanent, +-- systemic store error: every subsequent event would fail identically, so +-- dead-lettering would quietly drain the whole stream. +testDimensionMismatchHalts :: Assertion +testDimensionMismatchHalts = + withVectorEnv "dimension mismatch" \env capability -> do + decision <- + runHandler env capability (\_ -> pure (Right (Vector.replicate 8 0.1))) (Just 0) + case decision of + AckHalt (HaltFatal _) -> pure () + other -> assertFailure ("expected a fatal halt on a dimension mismatch, got: " <> show other) + +-- | Record a memory, then hand back its id and the @MemoryRecorded@ event at +-- the head of its stream — a real recorded event, not a hand-built one. +recordFixtureMemory :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + Text -> + Eff es (MemoryId, RecordedEvent) +recordFixtureMemory content = do + memoryId <- liftIO genMemoryId + now <- liftIO getCurrentTime + recorded <- + Memory.record + RecordMemoryData + { memoryId, + agentId = "test-agent", + sessionId = Nothing, + scope = fixtureScope, + memoryType = MemoryPreference, + content, + priority = 5, + confidence = HighConfidence, + tags = Set.fromList ["embedding-worker-spec"], + supersedes = Nothing, + recordedAt = now + } + void (liftIO (expectRight "Memory.record" recorded)) + events <- readStreamForward (Stream.streamName (memoryStream memoryId)) (StreamVersion 0) 10 + case Vector.toList events of + event : _ -> pure (memoryId, event) + [] -> liftIO (assertFailure "the recorded memory has no events") + +-- | Replace the event payload with something the memory codec cannot decode. +-- Spelled out field by field rather than as a record update because @payload@ +-- is also an 'Envelope' field, and GHC will not guess which one is meant. +corruptPayload :: RecordedEvent -> RecordedEvent +corruptPayload e = + RecordedEvent + { eventId = e.eventId, + eventType = e.eventType, + streamVersion = e.streamVersion, + globalPosition = e.globalPosition, + originalStreamId = e.originalStreamId, + originalVersion = e.originalVersion, + payload = Aeson.String "garbage", + metadata = e.metadata, + causationId = e.causationId, + correlationId = e.correlationId, + createdAt = e.createdAt + } + +runHandler :: + AppEnv -> + VectorCapability -> + (Text -> IO (Either EmbedError (Vector.Vector Double))) -> + Maybe Word -> + IO AckDecision +runHandler appEnv capability embed attemptN = + runOrFail appEnv do + (_, recorded) <- recordFixtureMemory "a memory to embed" + embeddingHandler capability (mkTestEnv embed) (mkIngested recorded attemptN) + +mkTestEnv :: (Text -> IO (Either EmbedError (Vector.Vector Double))) -> EmbeddingWorkerEnv +mkTestEnv embed = + EmbeddingWorkerEnv {model = testModel, dimensions = embeddingDims, embed} + +failingEmbed :: EmbedError -> (Text -> IO (Either EmbedError (Vector.Vector Double))) +failingEmbed err _ = pure (Left err) + +-- | The adapter's envelope, built by hand: only 'attempt' and 'payload' matter +-- to the handler, and the ack handle is never finalized because the handler +-- returns its decision rather than committing it. +mkIngested :: RecordedEvent -> Maybe Word -> Ingested es RecordedEvent +mkIngested recorded attemptN = + Ingested + { envelope = + Envelope + { messageId = MessageId "embedding-worker-spec", + cursor = Nothing, + partition = Nothing, + enqueuedAt = Nothing, + traceContext = Nothing, + headers = Nothing, + attempt = Attempt <$> attemptN, + attributes = HashMap.empty, + payload = recorded + }, + ack = AckHandle {finalize = \_ -> pure ()}, + lease = Nothing + } + +-- | Run a database-backed case, skipping it when the ephemeral PostgreSQL has +-- no pgvector: without the extension the embedding column does not exist, so +-- the store-error branches under test cannot be reached at all. +withVectorEnv :: String -> (AppEnv -> VectorCapability -> IO ()) -> Assertion +withVectorEnv label action = + withEmbeddingEnv \appEnv -> do + capability <- runOrFail appEnv (detectVectorCapability embeddingDims) + case capability of + VectorAvailable -> action appEnv capability + _ -> + putStrLn + ( " [skipped: " + <> label + <> "] pgvector is unavailable in the ephemeral database (" + <> show capability + <> ")" + ) + +withEmbeddingEnv :: (AppEnv -> IO a) -> IO a +withEmbeddingEnv action = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) action + +loadEmbeddingState :: (Store :> es) => Text -> Eff es Bool +loadEmbeddingState memoryId = + runTransaction (Tx.statement memoryId selectEmbeddedStmt) + +selectEmbeddedStmt :: Statement Text Bool +selectEmbeddedStmt = + preparable + """ + SELECT embedding IS NOT NULL AND content_hash IS NOT NULL + FROM kiroku.kioku_memories + WHERE memory_id = $1 + """ + (E.param (E.nonNullable E.text)) + (D.singleRow (D.column (D.nonNullable D.bool))) + +fixtureScope :: MemoryScope +fixtureScope = + ScopeEntity (Namespace "rei") (ScopeKind "intention") "intention_embedding_worker_spec" + +testModel :: EmbeddingModel +testModel = + toEmbeddingModel + EmbeddingConfig + { baseUrl = "http://embedding.invalid", + model = "text-embedding-3-small", + dimensions = embeddingDims, + apiKey = "not-used-the-provider-is-faked" + } + +-- | Matches the @vector(1536)@ column the embedding migration creates. +embeddingDims :: Int +embeddingDims = 1536 + +runOrFail :: AppEnv -> Eff AppEffects a -> IO a +runOrFail appEnv action = runAppIO appEnv action >>= expectRight "runAppIO" + +expectRight :: (Show e) => String -> Either e a -> IO a +expectRight label = \case + Left err -> assertFailure (label <> " failed: " <> show err) + Right value -> pure value diff --git a/test/Kioku/IdempotencySpec.hs b/test/Kioku/IdempotencySpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/IdempotencySpec.hs @@ -0,0 +1,421 @@ +-- | The idempotent-accept contract for session and memory writes. +-- +-- Every case here asserts two things: the right value came back, /and/ the event stream +-- gained no extra event. A write that returns @Right@ but quietly appended a second event +-- is not idempotent, and a write that returns a conflict but appended anyway has already +-- done the damage. +module Kioku.IdempotencySpec (tests) where + +import Control.Monad (void) +import Data.Set qualified as Set +import Data.Text (Text) +import Data.Time (UTCTime, getCurrentTime) +import Data.Vector qualified as Vector +import Effectful (Eff, IOE, liftIO, (:>)) +import Effectful.Error.Static (Error) +import Keiro.Stream qualified as Stream +import Kioku.Api.Scope (MemoryScope (..), Namespace (..)) +import Kioku.Api.Types (Confidence (..), MemoryType (..)) +import Kioku.App (AppEffects, runAppIO, withNoopAppEnv) +import Kioku.Id (MemoryId, SessionId, genMemoryId, genSessionId, idText) +import Kioku.Memory qualified as Memory +import Kioku.Memory.Domain (ArchiveMemoryData (..), RecordMemoryData (..), SupersedeMemoryData (..)) +import Kioku.Memory.EventStream (memoryStream) +import Kioku.Memory.ReadModel (MemoryRow (..)) +import Kioku.Migrations.TestSupport (withKiokuMigratedDatabase) +import Kioku.Session qualified as Session +import Kioku.Session.Domain + ( AwaitInputData (..), + CompleteSessionData (..), + FailSessionData (..), + ResumeSessionData (..), + StartSessionData (..), + ) +import Kioku.Session.EventStream (sessionStream) +import Kiroku.Store.Connection (defaultConnectionSettings) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Read (readStreamForward) +import Kiroku.Store.Types (StreamVersion (..)) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, assertFailure, testCase) + +tests :: TestTree +tests = + testGroup + "Idempotent accepts" + [ testGroup + "sessions" + [ testCase "an identical start is a duplicate" testStartDuplicate, + testCase "a start with a different focus is a conflict" testStartConflict, + testCase "an identical awaitInput is a duplicate" testAwaitDuplicate, + testCase "an awaitInput with a different reason is a conflict" testAwaitConflict, + testCase "an identical resume is a duplicate" testResumeDuplicate, + testCase "a resume with different input is a conflict" testResumeConflict, + testCase "an identical complete is a duplicate" testCompleteDuplicate, + testCase "completing a failed session is a conflict" testCompleteAfterFail, + testCase "failing a completed session is a conflict" testFailAfterComplete + ], + testGroup + "memories" + [ testCase "an identical record is a duplicate" testRecordDuplicate, + testCase "a record retried with a fresh clock is a duplicate" testRecordRetriedWithNewClock, + testCase "a record with different content is a conflict" testRecordConflict, + testCase "an identical supersede is a duplicate" testSupersedeDuplicate, + testCase "superseding by a different winner is a conflict" testSupersedeConflict, + testCase "archiving a superseded memory is a conflict" testArchiveAfterSupersede, + testCase "an identical merge is a duplicate" testMergeDuplicate, + testCase "merging into a different winner is a conflict" testMergeConflict + ] + ] + +-- * Sessions + +testStartDuplicate :: Assertion +testStartDuplicate = + withApp do + sid <- liftIO genSessionId + now <- liftIO getCurrentTime + let cmd = startData sid now + void (expectRight "first start" =<< Session.start cmd) + void (expectRight "duplicate start" =<< Session.start cmd) + assertSessionEvents sid 1 + +testStartConflict :: Assertion +testStartConflict = + withApp do + sid <- liftIO genSessionId + now <- liftIO getCurrentTime + let cmd = startData sid now + void (expectRight "first start" =<< Session.start cmd) + expectConflict "start with a different focus" + =<< Session.start cmd {focus = "a different focus"} + assertSessionEvents sid 1 + +testAwaitDuplicate :: Assertion +testAwaitDuplicate = + withApp do + sid <- startedSession + now <- liftIO getCurrentTime + let cmd = awaitData sid now + void (expectRight "first awaitInput" =<< Session.awaitInput cmd) + void (expectRight "duplicate awaitInput" =<< Session.awaitInput cmd) + assertSessionEvents sid 2 + +testAwaitConflict :: Assertion +testAwaitConflict = + withApp do + sid <- startedSession + now <- liftIO getCurrentTime + let cmd = awaitData sid now + void (expectRight "first awaitInput" =<< Session.awaitInput cmd) + expectConflict "awaitInput with a different reason" + =<< Session.awaitInput cmd {reason = "a different reason"} + assertSessionEvents sid 2 + +testResumeDuplicate :: Assertion +testResumeDuplicate = + withApp do + sid <- startedSession + now <- liftIO getCurrentTime + void (expectRight "awaitInput" =<< Session.awaitInput (awaitData sid now)) + let cmd = resumeData sid now "approved" + void (expectRight "first resume" =<< Session.resume cmd) + void (expectRight "duplicate resume" =<< Session.resume cmd) + assertSessionEvents sid 3 + +testResumeConflict :: Assertion +testResumeConflict = + withApp do + sid <- startedSession + now <- liftIO getCurrentTime + void (expectRight "awaitInput" =<< Session.awaitInput (awaitData sid now)) + void (expectRight "first resume" =<< Session.resume (resumeData sid now "approved")) + -- The session is running again; a re-delivery carrying a *different* answer is not this + -- request's own echo. + expectConflict "resume with different input" + =<< Session.resume (resumeData sid now "rejected") + assertSessionEvents sid 3 + +testCompleteDuplicate :: Assertion +testCompleteDuplicate = + withApp do + sid <- startedSession + now <- liftIO getCurrentTime + let cmd = completeData sid now + void (expectRight "first complete" =<< Session.complete cmd) + void (expectRight "duplicate complete" =<< Session.complete cmd) + assertSessionEvents sid 2 + +-- | The headline regression: this used to return @Right@ and report success for a session +-- that had actually failed. +testCompleteAfterFail :: Assertion +testCompleteAfterFail = + withApp do + sid <- startedSession + now <- liftIO getCurrentTime + void (expectRight "failSession" =<< Session.failSession (failData sid now)) + expectConflict "complete after fail" =<< Session.complete (completeData sid now) + assertSessionEvents sid 2 + +testFailAfterComplete :: Assertion +testFailAfterComplete = + withApp do + sid <- startedSession + now <- liftIO getCurrentTime + void (expectRight "complete" =<< Session.complete (completeData sid now)) + expectConflict "fail after complete" =<< Session.failSession (failData sid now) + assertSessionEvents sid 2 + +-- * Memories + +testRecordDuplicate :: Assertion +testRecordDuplicate = + withApp do + mid <- liftIO genMemoryId + now <- liftIO getCurrentTime + let cmd = recordData mid now "the original content" + void (expectRightM "first record" =<< Memory.record cmd) + void (expectRightM "duplicate record" =<< Memory.record cmd) + assertMemoryEvents mid 1 + +-- | @recordedAt@ must not participate in conflict detection. Distillation depends on this: +-- 'Kioku.Distill.L1.recordAtom' derives a deterministic memory id but stamps +-- @recordedAt = now@, so a timer re-fire re-records the same atom with a later clock. If +-- that were a conflict, every re-fire of an L1 pass would hard-fail — which is exactly what +-- happened when this contract was first written the other way. +testRecordRetriedWithNewClock :: Assertion +testRecordRetriedWithNewClock = + withApp do + mid <- liftIO genMemoryId + firstAt <- liftIO getCurrentTime + let content = "identical content, later clock" + void (expectRightM "first record" =<< Memory.record (recordData mid firstAt content)) + laterAt <- liftIO getCurrentTime + void (expectRightM "retry with a fresh clock" =<< Memory.record (recordData mid laterAt content)) + assertMemoryEvents mid 1 + lookedUp <- Memory.getMemoryRowById mid + liftIO case lookedUp of + Left err -> assertFailure ("lookup: " <> show err) + Right Nothing -> assertFailure "the memory row vanished" + Right (Just r) -> assertEqual "the first write's createdAt is kept" firstAt r.createdAt + +testRecordConflict :: Assertion +testRecordConflict = + withApp do + mid <- liftIO genMemoryId + now <- liftIO getCurrentTime + void (expectRightM "first record" =<< Memory.record (recordData mid now "the original content")) + expectConflictM "record with different content" + =<< Memory.record (recordData mid now "something else entirely") + assertMemoryEvents mid 1 + +testSupersedeDuplicate :: Assertion +testSupersedeDuplicate = + withApp do + loser <- recordedMemory "loser" + winner <- recordedMemory "winner" + now <- liftIO getCurrentTime + let cmd = SupersedeMemoryData {memoryId = loser, supersededBy = winner, supersededAt = now} + void (expectRightM "first supersede" =<< Memory.supersede cmd) + void (expectRightM "duplicate supersede" =<< Memory.supersede cmd) + assertMemoryEvents loser 2 + +-- | The other headline regression: supersede by X, then by Y, used to report success for Y +-- while X remained the recorded winner. +testSupersedeConflict :: Assertion +testSupersedeConflict = + withApp do + loser <- recordedMemory "loser" + winnerX <- recordedMemory "winner x" + winnerY <- recordedMemory "winner y" + now <- liftIO getCurrentTime + void $ + expectRightM "supersede by X" + =<< Memory.supersede SupersedeMemoryData {memoryId = loser, supersededBy = winnerX, supersededAt = now} + expectConflictM "supersede by Y after X" + =<< Memory.supersede SupersedeMemoryData {memoryId = loser, supersededBy = winnerY, supersededAt = now} + assertMemoryEvents loser 2 + +testArchiveAfterSupersede :: Assertion +testArchiveAfterSupersede = + withApp do + loser <- recordedMemory "loser" + winner <- recordedMemory "winner" + now <- liftIO getCurrentTime + void $ + expectRightM "supersede" + =<< Memory.supersede SupersedeMemoryData {memoryId = loser, supersededBy = winner, supersededAt = now} + expectConflictM "archive after supersede" + =<< Memory.archive ArchiveMemoryData {memoryId = loser, archivedAt = now} + assertMemoryEvents loser 2 + +testMergeDuplicate :: Assertion +testMergeDuplicate = + withApp do + loser <- recordedMemory "loser" + winner <- recordedMemory "winner" + void (expectRightM "first merge" =<< Memory.merge loser winner) + void (expectRightM "duplicate merge" =<< Memory.merge loser winner) + assertMemoryEvents loser 2 + +testMergeConflict :: Assertion +testMergeConflict = + withApp do + loser <- recordedMemory "loser" + winnerX <- recordedMemory "winner x" + winnerY <- recordedMemory "winner y" + void (expectRightM "merge into X" =<< Memory.merge loser winnerX) + expectConflictM "merge into Y after X" =<< Memory.merge loser winnerY + assertMemoryEvents loser 2 + +-- * Fixtures + +startData :: SessionId -> UTCTime -> StartSessionData +startData sid startedAt = + StartSessionData + { sessionId = sid, + agentId = "test-agent", + focus = "idempotency", + scope = testScope, + subjectRef = Nothing, + previousSessionId = Nothing, + parentSessionId = Nothing, + delegationDepth = 0, + startedAt + } + +awaitData :: SessionId -> UTCTime -> AwaitInputData +awaitData sid awaitedAt = + AwaitInputData + { sessionId = sid, + reason = "approval", + correlationKey = Just "k1", + deadline = Nothing, + awaitedAt + } + +resumeData :: SessionId -> UTCTime -> Text -> ResumeSessionData +resumeData sid resumedAt input = + ResumeSessionData + { sessionId = sid, + correlationKey = Just "k1", + force = False, + input, + resumedAt + } + +completeData :: SessionId -> UTCTime -> CompleteSessionData +completeData sid completedAt = + CompleteSessionData + { sessionId = sid, + completedAt, + modelUsed = Just "test-model", + summary = Just "done" + } + +failData :: SessionId -> UTCTime -> FailSessionData +failData sid failedAt = + FailSessionData + { sessionId = sid, + failedAt, + errorMessage = "boom" + } + +recordData :: MemoryId -> UTCTime -> Text -> RecordMemoryData +recordData mid recordedAt content = + RecordMemoryData + { memoryId = mid, + agentId = "test-agent", + sessionId = Nothing, + scope = testScope, + memoryType = MemoryFact, + content, + priority = 50, + confidence = HighConfidence, + tags = Set.fromList ["t"], + supersedes = Nothing, + recordedAt + } + +startedSession :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + Eff es SessionId +startedSession = do + sid <- liftIO genSessionId + now <- liftIO getCurrentTime + void (expectRight "Session.start" =<< Session.start (startData sid now)) + pure sid + +recordedMemory :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + Text -> + Eff es MemoryId +recordedMemory content = do + mid <- liftIO genMemoryId + now <- liftIO getCurrentTime + void (expectRightM "Memory.record" =<< Memory.record (recordData mid now content)) + pure mid + +-- * Assertions + +assertSessionEvents :: + (IOE :> es, Store :> es) => + SessionId -> + Int -> + Eff es () +assertSessionEvents sid expected = do + events <- readStreamForward (Stream.streamName (sessionStream sid)) (StreamVersion 0) 100 + liftIO $ + assertEqual + "the session stream gained no extra event" + expected + (Vector.length events) + +assertMemoryEvents :: + (IOE :> es, Store :> es) => + MemoryId -> + Int -> + Eff es () +assertMemoryEvents mid expected = do + events <- readStreamForward (Stream.streamName (memoryStream mid)) (StreamVersion 0) 100 + liftIO $ + assertEqual + "the memory stream gained no extra event" + expected + (Vector.length events) + +expectRight :: (IOE :> es) => String -> Either Session.SessionWriteError a -> Eff es a +expectRight label = \case + Right value -> pure value + Left err -> liftIO (assertFailure (label <> ": expected success, got " <> show err)) + +expectConflict :: (IOE :> es) => String -> Either Session.SessionWriteError SessionId -> Eff es () +expectConflict label = \case + Left (Session.SessionConflict reason) -> + liftIO $ assertBool (label <> ": the conflict names a field") (reason /= "") + other -> liftIO (assertFailure (label <> ": expected SessionConflict, got " <> show other)) + +expectRightM :: (IOE :> es) => String -> Either Memory.MemoryWriteError a -> Eff es a +expectRightM label = \case + Right value -> pure value + Left err -> liftIO (assertFailure (label <> ": expected success, got " <> show err)) + +expectConflictM :: (IOE :> es) => String -> Either Memory.MemoryWriteError MemoryId -> Eff es () +expectConflictM label = \case + Left (Memory.MemoryConflict reason) -> + liftIO $ assertBool (label <> ": the conflict names a field") (reason /= "") + other -> liftIO (assertFailure (label <> ": expected MemoryConflict, got " <> show other)) + +withApp :: Eff AppEffects a -> IO a +withApp action = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) \env -> do + result <- runAppIO env action + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right value -> pure value + +testScope :: MemoryScope +testScope = ScopeGlobal (Namespace "kioku-test") diff --git a/test/Kioku/ReadModelReconcileSpec.hs b/test/Kioku/ReadModelReconcileSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/ReadModelReconcileSpec.hs @@ -0,0 +1,168 @@ +-- | The fail-closed outage and its repair, end to end. +-- +-- Keiro refuses to serve a read model whose registry row disagrees with the code's +-- declared version — a deliberate safety property. The hazard is that nothing used to +-- bring those rows back into agreement: 'Keiro.ReadModel.Schema.registerReadModel' only +-- inserts, so an additive migration that advances a model's version leaves every query +-- for it failing with 'ReadModelStaleSchema' until a human hand-writes a registry +-- migration. That is not hypothetical; it is what happened when the session models went +-- v1 -> v2 -> v3. +-- +-- 'reconcileReadModelRegistry' is the repair, and this spec walks the whole arc: startup +-- registration, a healthy query, a downgraded registry row, the resulting outage, the +-- reconcile, and the query working again. +module Kioku.ReadModelReconcileSpec (tests) where + +import Data.Text qualified as Text +import Data.Text.Encoding (encodeUtf8) +import Data.Time (getCurrentTime) +import Effectful (Eff, IOE, liftIO, (:>)) +import Effectful.Error.Static (Error) +import Hasql.Transaction qualified as Tx +import Keiro.ReadModel (ReadModelError (..)) +import Kioku.Api.Scope (MemoryScope (..), Namespace (..)) +import Kioku.App (AppEffects, runAppIO, withNoopAppEnv) +import Kioku.Id (SessionId, genSessionId) +import Kioku.Migrations.TestSupport (withKiokuMigratedDatabase) +import Kioku.Prelude +import Kioku.ReadModel + ( ReadModelSchema (..), + ReconcileOutcome (..), + kiokuReadModelSchemas, + reconcileReadModelRegistry, + ) +import Kioku.Session qualified as Session +import Kioku.Session.Domain (StartSessionData (..)) +import Kiroku.Store.Connection (defaultConnectionSettings) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Transaction (runTransaction) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, assertFailure, testCase) + +tests :: TestTree +tests = + testGroup + "ReadModel.Reconcile" + [ testCase "a stale registry row fails every query closed, then reconciles" testStaleThenReconcile, + testCase "reconciliation is idempotent" testIdempotent + ] + +-- | The whole arc. Each step is asserted, including the outage itself — without that +-- assertion the test could pass against a build where the guard never fires at all. +testStaleThenReconcile :: Assertion +testStaleThenReconcile = + withApp \sid -> do + -- Application startup registered every model at the identity the code declares. + healthy <- Session.getById sid + liftIO $ assertBool "a fresh database serves session queries" (isRight healthy) + + downgradeSessionByIdTo 2 "kioku-session-v2" + + stale <- Session.getById sid + liftIO case stale of + Left (ReadModelStaleSchema name expectedVersion foundVersion expectedHash foundHash) -> do + assertEqual "the stale model" "kioku-session-by-id" name + assertEqual "expected version" 3 expectedVersion + assertEqual "found version" 2 foundVersion + assertEqual "expected hash" "kioku-session-v3" expectedHash + assertEqual "found hash" "kioku-session-v2" foundHash + other -> + assertFailure + ("expected the query to fail closed on the stale row, got " <> show (() <$ other)) + + outcomes <- reconcileReadModelRegistry + liftIO do + assertEqual + "the downgraded row was bumped back" + (Just Reconciled) + (outcomeFor "kioku-session-by-id" outcomes) + -- Reconciliation must cover every model the code declares, not just the one this + -- test downgraded. Startup registered the rest, so they stay current. + assertEqual + "every declared model was accounted for" + (map (.readModelName) kiokuReadModelSchemas) + (map ((.readModelName) . fst) outcomes) + assertEqual + "the models that were not downgraded stayed current" + [] + [ schema.readModelName + | (schema, outcome) <- outcomes, + schema.readModelName /= "kioku-session-by-id", + outcome /= AlreadyCurrent + ] + + repaired <- Session.getById sid + liftIO $ assertBool "the query works again" (isRight repaired) + +-- | A second pass must write nothing. If it reported 'Reconciled' again, the reconciler +-- would be rewriting @last_built_at@ on every @just migrate@ — and, worse, would be lying +-- about what it changed. +testIdempotent :: Assertion +testIdempotent = + withApp \sid -> do + void (Session.getById sid) + _ <- reconcileReadModelRegistry + second <- reconcileReadModelRegistry + liftIO $ + assertEqual + "every model is already current on the second pass" + [] + [schema.readModelName | (schema, outcome) <- second, outcome /= AlreadyCurrent] + +-- | Pin the registry row back to an older identity, exactly as a database that missed the +-- v3 bump would have it. The name is unqualified so it resolves through the store's +-- @search_path@, precisely as keiro's own registry statements do. +downgradeSessionByIdTo :: (Store :> es) => Int -> Text -> Eff es () +downgradeSessionByIdTo version shapeHash = + runTransaction . Tx.sql . encodeUtf8 $ + "UPDATE keiro.keiro_read_models SET version = " + <> Text.pack (show version) + <> ", shape_hash = '" + <> shapeHash + <> "' WHERE name = 'kioku-session-by-id'" + +outcomeFor :: Text -> [(ReadModelSchema, ReconcileOutcome)] -> Maybe ReconcileOutcome +outcomeFor name outcomes = + lookup name [(schema.readModelName, outcome) | (schema, outcome) <- outcomes] + +isRight :: Either e a -> Bool +isRight = \case + Right _ -> True + Left _ -> False + +-- * Harness + +withApp :: (SessionId -> Eff AppEffects ()) -> Assertion +withApp use = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) \env -> do + sid <- genSessionId + result <- runAppIO env (startFixture sid >> use sid) + case result of + Left err -> assertFailure ("store error: " <> show err) + Right () -> pure () + +startFixture :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + Eff es () +startFixture sid = do + now <- liftIO getCurrentTime + result <- + Session.start + StartSessionData + { sessionId = sid, + agentId = "test-agent", + focus = "read-model reconciliation", + scope = ScopeGlobal (Namespace "kioku-test"), + subjectRef = Nothing, + previousSessionId = Nothing, + parentSessionId = Nothing, + delegationDepth = 0, + startedAt = now + } + case result of + Left err -> liftIO (assertFailure ("Session.start: " <> show err)) + Right _ -> pure () diff --git a/test/Kioku/RecallHarness.hs b/test/Kioku/RecallHarness.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/RecallHarness.hs @@ -0,0 +1,522 @@ +-- | An instrument for measuring the quality of kioku's vector recall against a corpus whose +-- true answer is known /by construction/, in Haskell, without asking Postgres anything. +-- +-- == Why this exists +-- +-- Recall's vector channel can silently return nothing. The HNSW index covers the embedding +-- column alone (its only predicate is @embedding IS NOT NULL@), so it picks its candidates by +-- distance and the namespace, scope, and @status = 'active'@ predicates are applied /afterwards/, +-- to rows the index has already chosen. When the rows outside the caller's scope are nearer the +-- query than the in-scope answers, the index can spend its whole budget on rows the filter then +-- discards. That is /filtered-ANN starvation/, and because 'Kioku.Recall.fuseRecallCandidates' +-- blends the two channels by rank, an empty vector channel simply contributes no ranks: the +-- score degrades smoothly into pure keyword scoring, with no error and no warning. Nothing in a +-- 'Kioku.Api.Types.MemoryRecord' or a @RecallHit@ says "the semantic half came back empty". +-- +-- This module builds the corpus that provokes that, and measures what actually came back. +-- +-- == The geometry, and why it is the point +-- +-- The query sits on axis 0. A seeded vector at angle @t@ (radians) is @cos t * e0 + sin t * e1@, +-- so its cosine distance to the query is exactly @1 - cos t@ — a pure, monotonically increasing +-- function of one knob on @[0, pi]@. The harness therefore knows the true ranking of every +-- seeded row /without querying the database/, which is what makes recall@k a measurement rather +-- than a tautology. The obvious alternative — run an exact scan and compare the approximate +-- result to it — would measure the ANN path against the planner's /other/ choice, and the +-- planner's choice between those two plans is precisely the thing under suspicion. +-- +-- (The existing 'Kioku.RecallSqlSpec.unitVector' cannot express this: orthogonal basis vectors +-- sit at cosine distance exactly 0 or exactly 1, and starvation needs a graded scale.) +-- +-- == Two traps that produce confidently-wrong numbers rather than errors +-- +-- 1. __Rows inserted inside an open transaction never get an HNSW index scan__, not even with +-- @enable_seqscan = off@ and a fresh @ANALYZE@. Every seeding statement here is its own +-- committed transaction ('runTransaction' commits). Do not "optimise" 'seedCorpus' by +-- wrapping the whole corpus in one transaction — the rows would be invisible to the index +-- and every number downstream would be fiction. +-- 2. __A partial index needs its predicate restated in the query__, or the planner cannot prove +-- the index applies and falls back to a sequential scan. The @embedding IS NOT NULL@ in +-- 'explainVectorStmt' is load-bearing, not decoration. +-- +-- See docs/plans/18-build-a-recall-quality-harness-that-reproduces-filtered-ann-starvation.md. +module Kioku.RecallHarness + ( -- * Geometry + vectorAtAngle, + queryVector, + cosineDistanceAtAngle, + embeddingDimensions, + + -- * Seeding + CorpusConfig (..), + defaultStarvationCorpus, + SeededCorpus (..), + seedCorpus, + + -- * Measurement + RecallQuality (..), + measureRecallQuality, + measureRecallQualityWith, + explainVectorQuery, + explainVectorQueryWith, + describeRecallQuality, + usedHnswIndex, + planAgreesWithQuery, + runDdl, + ) +where + +import Data.Char (isDigit) +import Data.Foldable (traverse_) +import Data.Functor.Contravariant ((>$<)) +import Data.Int (Int32) +import Data.List (isInfixOf, sortOn) +import Data.Set qualified as Set +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding (encodeUtf8) +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import Effectful (Eff, (:>)) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), scopeKindText, scopeNamespaceText, scopeRefText) +import Kioku.Api.Types (MemoryRecord (..)) +import Kioku.Recall + ( RecallRequest (..), + RecallStrategy (..), + VectorChannelOutcome (..), + candidatePoolSize, + memoryRecordColumns, + selectVectorCandidatesDiagnosed, + selectVectorCandidatesStmt, + vectorCandidateQuery, + ) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Transaction (runTransaction) +import Text.Read (readMaybe) + +-- * Geometry + +-- | The width of @kioku_memories.embedding@, which is @vector(1536)@. A seeded vector must +-- match it exactly or the @::vector@ cast fails. +embeddingDimensions :: Int +embeddingDimensions = 1536 + +-- | @cos t * e0 + sin t * e1@ — a unit vector at angle @t@ from the query axis, in the plane +-- spanned by the first two coordinates. Every other coordinate is zero. +vectorAtAngle :: Double -> Vector Double +vectorAtAngle t = + Vector.generate embeddingDimensions \j -> + case j of + 0 -> cos t + 1 -> sin t + _ -> 0 + +-- | The query the whole harness measures against: the vector at angle zero, i.e. @e0@. +queryVector :: Vector Double +queryVector = vectorAtAngle 0 + +-- | The cosine distance from 'queryVector' to 'vectorAtAngle', in closed form. +-- +-- Both vectors are unit length, so cosine distance is @1 - cos(angle between them)@, and the +-- angle between @e0@ and @vectorAtAngle t@ is @t@. This is the ground truth: it is computed +-- here, in Haskell, and never read back from the database, which is the thing under test. +cosineDistanceAtAngle :: Double -> Double +cosineDistanceAtAngle t = 1 - cos t + +-- * Seeding + +-- | The knobs that make a corpus starve. +-- +-- Starvation needs the filter to /correlate with distance/: the rows the scope filter throws +-- away must be the ones the index reaches for first. So the decoys live in a different +-- namespace and sit /nearer/ the query than any in-scope row. +data CorpusConfig = CorpusConfig + { -- | Memories in the namespace and scope the query asks for. These are the true answers. + inScopeCount :: !Int, + -- | Memories in a /different/ namespace, which the query must never return. + decoyCount :: !Int, + -- | The angular band (radians) the in-scope rows are spread evenly across. + inScopeAngles :: !(Double, Double), + -- | The angular band the decoys occupy. Make it strictly nearer the query than + -- 'inScopeAngles' — that is, smaller angles — or nothing starves. + decoyAngles :: !(Double, Double) + } + deriving stock (Eq, Show) + +-- | The probe the previous initiative recorded as "1648 rows removed by filter, zero returned": +-- 2000 in-scope memories, 2000 nearer decoys in another namespace. +-- +-- Every decoy is strictly nearer the query than every in-scope row. In cosine distance the +-- decoys span roughly 0.001 to 0.12 and the in-scope rows roughly 0.30 to 0.64, so the index, +-- descending towards the query, meets all 2000 decoys before the first true answer. +defaultStarvationCorpus :: CorpusConfig +defaultStarvationCorpus = + CorpusConfig + { inScopeCount = 2000, + decoyCount = 2000, + inScopeAngles = (0.8, 1.2), + decoyAngles = (0.05, 0.5) + } + +-- | What was seeded, including the ground truth. +data SeededCorpus = SeededCorpus + { config :: !CorpusConfig, + -- | The scope the query asks for. Its namespace holds only the in-scope rows. + targetScope :: !MemoryScope, + -- | The in-scope memory ids ordered by true cosine distance, nearest first. Computed from + -- the seed angles, never read back from the database. + trueNearestInScope :: ![Text] + } + deriving stock (Eq, Show) + +-- | The namespace the query asks for. Only in-scope rows live here. +targetNamespace :: Text +targetNamespace = "harness_target" + +-- | The namespace the decoys live in. The query must never return one of these. +decoyNamespace :: Text +decoyNamespace = "harness_decoy" + +-- | Spread @n@ points evenly across @[lo, hi]@, inclusive at both ends. +anglesAcross :: Int -> (Double, Double) -> [Double] +anglesAcross n (lo, hi) + | n <= 0 = [] + | n == 1 = [lo] + | otherwise = + [ lo + (hi - lo) * fromIntegral i / fromIntegral (n - 1) + | i <- [0 .. n - 1] + ] + +-- | Seed the corpus and return its ground truth. +-- +-- Rows are inserted in committed batches (see the module header's trap 1: rows in an open +-- transaction get no index scan), and the table is @ANALYZE@d afterwards, because without +-- statistics the planner uses defaults and the plan it picks is not the plan production would +-- pick — which, for this harness, is the entire subject. +seedCorpus :: (Store :> es) => CorpusConfig -> Eff es SeededCorpus +seedCorpus cfg = do + let inScope = + [ (inScopeId i, targetNamespace, t) + | (i, t) <- zip [0 :: Int ..] (anglesAcross cfg.inScopeCount cfg.inScopeAngles) + ] + decoys = + [ (decoyId i, decoyNamespace, t) + | (i, t) <- zip [0 :: Int ..] (anglesAcross cfg.decoyCount cfg.decoyAngles) + ] + traverse_ insertBatch (chunksOf seedBatchSize (inScope <> decoys)) + runTransaction (Tx.sql "ANALYZE kioku_memories") + pure + SeededCorpus + { config = cfg, + targetScope = ScopeGlobal (Namespace targetNamespace), + -- Distance is @1 - cos t@, which increases monotonically with @t@ on @[0, pi]@, so + -- ordering by angle *is* ordering by distance. 'anglesAcross' already emits ascending + -- angles; sorting explicitly keeps that from being a silent assumption. + trueNearestInScope = + fmap (\(memoryId, _, _) -> memoryId) (sortOn (\(_, _, t) -> t) inScope) + } + where + inScopeId i = "m_in_" <> Text.pack (show i) + decoyId i = "m_decoy_" <> Text.pack (show i) + +-- | Rows per @INSERT@. Each batch is its own committed transaction. +seedBatchSize :: Int +seedBatchSize = 500 + +insertBatch :: (Store :> es) => [(Text, Text, Double)] -> Eff es () +insertBatch [] = pure () +insertBatch rows = + runTransaction . Tx.sql . encodeUtf8 $ + "INSERT INTO kioku_memories \ + \(memory_id, agent_id, namespace, scope_kind, scope_ref, memory_type, content, status, created_at, updated_at, embedding) VALUES " + <> Text.intercalate ", " (row <$> rows) + where + row (memoryId, namespace, t) = + "('" + <> memoryId + <> "', 'agent', '" + <> namespace + <> "', NULL, NULL, 'fact', 'seeded corpus row " + <> memoryId + <> "', 'active', now(), now(), " + <> sparseVectorSql t + <> ")" + +-- | The seeded vector, built as SQL rather than as a 1536-element text literal. +-- +-- 'Kioku.Recall.vectorLiteral' would render all 1536 components, of which 1534 are zero: about +-- 6KB per row, so around 25MB of SQL text for the default corpus and far more for the sweep. +-- Only the first two coordinates are non-zero, so the zeros are appended by Postgres with +-- @repeat@ instead. The two significant components are still computed in Haskell — the ground +-- truth stays Haskell's — and the M1 instrument case asserts that the distances Postgres +-- actually computes for these rows match @1 - cos t@, which is what makes the shortcut safe +-- rather than merely clever. +sparseVectorSql :: Double -> Text +sparseVectorSql t = + "('[" + <> showDouble (cos t) + <> "," + <> showDouble (sin t) + <> "' || repeat(',0', " + <> Text.pack (show (embeddingDimensions - 2)) + <> ") || ']')::vector" + +showDouble :: Double -> Text +showDouble = Text.pack . show + +chunksOf :: Int -> [a] -> [[a]] +chunksOf n xs + | n <= 0 = [xs] + | otherwise = case splitAt n xs of + (chunk, []) -> [chunk | not (null chunk)] + (chunk, rest) -> chunk : chunksOf n rest + +-- * Measurement + +-- | What the vector channel actually did. +data RecallQuality = RecallQuality + { -- | How many candidates the vector channel produced. The pool is 50 + -- ('Kioku.Recall.candidatePoolSize'), so a healthy selective scope returns 50. + rowsReturned :: !Int, + -- | Of the @k@ truly nearest in-scope memories, what fraction came back. 1.0 is perfect; + -- 0.0 means the search found none of them. + recallAtK :: !Double, + k :: !Int, + -- | How many of the returned rows were decoys. Must always be zero: the scope filter is a + -- correctness boundary, and a non-zero value here means something far worse than starvation. + decoysReturned :: !Int, + -- | Rows the approximate (HNSW) pass returned, before any exact fallback. + annRows :: !Int, + -- | Whether the exact fallback ran because the approximate pass came back short. + exactFallbackFired :: !Bool, + -- | @EXPLAIN (ANALYZE, BUFFERS)@ for the query as it is actually issued. Carried so a + -- failing case can print the cause rather than just @expected: True, got: False@. + planText :: !Text, + -- | The row count the captured plan's top node actually produced. If this disagrees with + -- 'rowsReturned', the EXPLAIN is describing a query nobody runs — see 'planAgreesWithQuery'. + planTopRows :: !(Maybe Int) + } + deriving stock (Eq, Show) + +-- | Does the captured plan describe the query that was actually measured? +-- +-- This is the instrument's self-check, and it exists because the instrument got this wrong +-- once. An @EXPLAIN@ whose SQL differs from the real statement in any way the planner cares +-- about — notably the width of the select list, which sets the cost of the top-N sort the +-- /exact/ plan needs — can choose a different plan and report a different result. The failure +-- mode is silent and flattering: the EXPLAIN says "50 rows, all good" while the real query +-- returns zero. +-- +-- So: the plan's top node must have produced the same number of rows the query returned. If it +-- did not, every conclusion drawn from 'planText' is void, and a case built on it must fail +-- loudly rather than report a comfortable number. +-- +-- The comparison is against 'annRows', not 'rowsReturned', because 'planText' captures the +-- /approximate/ statement, and the channel may then run an exact fallback whose rows the ANN plan +-- naturally does not account for. Comparing against the final count would make this check fail +-- precisely when the fallback is doing its job — which would be a false alarm, and a false alarm +-- that fires every time is a check nobody keeps. +planAgreesWithQuery :: RecallQuality -> Bool +planAgreesWithQuery q = maybe False (== q.annRows) q.planTopRows + +-- | Pull @rows=N@ out of the @(actual time=… rows=N loops=1)@ segment of the plan's first line, +-- which is its top node. Returns 'Nothing' if the plan has no @actual@ section, which would +-- mean the EXPLAIN ran without @ANALYZE@ and is not a measurement at all. +planActualTopRows :: Text -> Maybe Int +planActualTopRows plan = do + firstLine <- case Text.lines plan of + l : _ -> Just l + [] -> Nothing + actual <- afterToken "actual" firstLine + rows <- afterToken "rows=" actual + readMaybe (Text.unpack (Text.takeWhile isDigit rows)) + where + afterToken token haystack = + case Text.breakOn token haystack of + (_, rest) + | Text.null rest -> Nothing + | otherwise -> Just (Text.drop (Text.length token) rest) + +-- | Run the vector channel against a seeded corpus and score what it returned. +-- +-- This drives 'Kioku.Recall.selectVectorCandidates' — the exported test seam — rather than a +-- copy of the SQL, so when the statement changes the harness measures the /new/ one and cannot +-- silently keep testing a query that no longer runs in production. +measureRecallQuality :: (Store :> es) => SeededCorpus -> Int -> Eff es RecallQuality +measureRecallQuality corpus k = do + (outcome, rows) <- selectVectorCandidatesDiagnosed (vectorRequest corpus) queryVector + plan <- explainVectorQuery corpus + pure (scoreRecallQuality corpus k rows plan outcome.annRows outcome.exactFallbackFired) + +-- | 'measureRecallQuality', but with @SET LOCAL@ settings applied to the transaction the vector +-- statement runs in — the seam the bake-off in +-- docs/plans/19-fix-filtered-ann-starvation-in-vector-recall.md needs. +-- +-- Two details make this a real measurement rather than a plausible one. +-- +-- First, the settings and the query must share a transaction. @SET LOCAL@ lasts exactly as long +-- as the surrounding transaction, so issuing it in one @runTransaction@ and the query in another +-- is a no-op that silently measures the baseline while claiming to measure the candidate — and +-- reports a confident number either way. Both go in the single transaction below. +-- +-- Second, the query is 'Kioku.Recall.selectVectorCandidatesStmt' itself, not a copy, and the +-- @EXPLAIN@ runs under the same settings in its own transaction. 'planAgreesWithQuery' still +-- guards the pair. +measureRecallQualityWith :: + (Store :> es) => + -- | @SET LOCAL@ statements, e.g. @["SET LOCAL hnsw.iterative_scan = 'strict_order'"]@. + [Text] -> + SeededCorpus -> + Int -> + Eff es RecallQuality +measureRecallQualityWith settings corpus k = do + rows <- runTransaction do + traverse_ (Tx.sql . encodeUtf8) settings + Tx.statement (vectorCandidateQuery (vectorRequest corpus) queryVector) selectVectorCandidatesStmt + plan <- explainVectorQueryWith settings corpus + -- This drives the raw ANN statement, with no fallback, so the ANN pass *is* the whole channel. + pure (scoreRecallQuality corpus k rows plan (length rows) False) + +scoreRecallQuality :: SeededCorpus -> Int -> [MemoryRecord] -> Text -> Int -> Bool -> RecallQuality +scoreRecallQuality corpus k rows plan annPassRows fallbackFired = + RecallQuality + { rowsReturned = length rows, + recallAtK = + if null truth + then 0 + else fromIntegral found / fromIntegral (length truth), + k, + decoysReturned = length (filter ("m_decoy_" `Text.isPrefixOf`) returned), + annRows = annPassRows, + exactFallbackFired = fallbackFired, + planText = plan, + planTopRows = planActualTopRows plan + } + where + returned = (\row -> row.memoryId) <$> rows + returnedSet = Set.fromList returned + truth = take k corpus.trueNearestInScope + found = length (filter (`Set.member` returnedSet) truth) + +-- | The request the vector channel is driven with. The query /text/ is irrelevant — the vector +-- statement never reads it — but 'RecallRequest' requires one. +vectorRequest :: SeededCorpus -> RecallRequest +vectorRequest corpus = + RecallRequest + { scope = corpus.targetScope, + query = "seeded corpus row", + strategy = Embedding, + maxResults = 10 + } + +-- | @EXPLAIN (ANALYZE, BUFFERS)@ for the vector candidate query, run with the same five +-- parameters recall passes, so Postgres plans it the way it plans the real one. +explainVectorQuery :: (Store :> es) => SeededCorpus -> Eff es Text +explainVectorQuery = explainVectorQueryWith [] + +-- | 'explainVectorQuery' with @SET LOCAL@ settings applied to the same transaction. They must +-- share a transaction: @SET LOCAL@ lasts exactly as long as the surrounding one, so a setting +-- issued separately would silently EXPLAIN the baseline while claiming to EXPLAIN the candidate. +explainVectorQueryWith :: (Store :> es) => [Text] -> SeededCorpus -> Eff es Text +explainVectorQueryWith settings corpus = + Text.unlines <$> runTransaction do + traverse_ (Tx.sql . encodeUtf8) settings + Tx.statement + ( Text.pack (show (Vector.toList queryVector)), + scopeNamespaceText corpus.targetScope, + scopeKindText corpus.targetScope, + scopeRefText corpus.targetScope, + candidatePoolSize + ) + explainVectorStmt + +-- | Everything here is copied verbatim from 'Kioku.Recall.selectVectorCandidatesStmt' — the +-- select list, the predicates, the @ORDER BY@, and the @LIMIT@ — and every part of it earns +-- its place. +-- +-- @embedding IS NOT NULL@ is not decoration: the HNSW index is partial on exactly that +-- predicate, and without it restated here the planner cannot prove the index applies and falls +-- back to a sequential scan — which would look like "the index is broken" and be entirely an +-- artifact of the measurement. +-- +-- __The select list is load-bearing too, which is not obvious and was learned the hard way.__ +-- An earlier version of this function selected @memory_id@ alone, on the theory that Postgres +-- chooses the plan from the @WHERE@, the @ORDER BY@ and the @LIMIT@, and that the projection +-- could not turn an HNSW scan into anything else. That is false. The projection sets the row +-- width, the width sets the cost of the top-N sort that the /exact/ plan needs, and that cost +-- is exactly what the planner weighs against the HNSW scan. On the 2000-in-scope, 2000-decoy +-- corpus the narrow projection made the sort look cheap, the planner took the exact plan, and +-- the EXPLAIN reported 50 happy rows — while the real query, with its 13 real columns, took +-- the HNSW plan and returned zero. The instrument was describing a query nobody runs. +explainVectorStmt :: Statement (Text, Text, Maybe Text, Maybe Text, Int32) [Text] +explainVectorStmt = + preparable + ( "EXPLAIN (ANALYZE, BUFFERS) SELECT " + <> memoryRecordColumns + <> " FROM kiroku.kioku_memories \ + \ WHERE status = 'active' \ + \ AND namespace = $2 \ + \ AND (($3 IS NULL AND $4 IS NULL) OR (scope_kind = $3 AND scope_ref = $4)) \ + \ AND embedding IS NOT NULL \ + \ ORDER BY embedding <=> $1::vector \ + \ LIMIT $5" + ) + encoder + (D.rowList (D.column (D.nonNullable D.text))) + where + encoder = + ((\(v, _, _, _, _) -> v) >$< E.param (E.nonNullable E.text)) + <> ((\(_, n, _, _, _) -> n) >$< E.param (E.nonNullable E.text)) + <> ((\(_, _, sk, _, _) -> sk) >$< E.param (E.nullable E.text)) + <> ((\(_, _, _, sr, _) -> sr) >$< E.param (E.nullable E.text)) + <> ((\(_, _, _, _, l) -> l) >$< E.param (E.nonNullable E.int4)) + +-- | A failure message that reads as a diagnosis rather than an assertion. +-- +-- The whole purpose of this harness is to convert an invisible failure into a legible one, so a +-- case that fails with @expected: True, got: False@ has not delivered it. The most telling line +-- in the plan is @Rows Removed by Filter@ — starvation, made visible. +describeRecallQuality :: RecallQuality -> String +describeRecallQuality q = + Text.unpack . Text.unlines $ + [ "the vector channel returned " + <> Text.pack (show q.rowsReturned) + <> " candidates, recall@" + <> Text.pack (show q.k) + <> " = " + <> Text.pack (show q.recallAtK) + <> (if q.decoysReturned > 0 then " (!! " <> Text.pack (show q.decoysReturned) <> " out-of-scope decoys leaked)" else ""), + "plan: " + <> (if usedHnswIndex q then "HNSW (approximate)" else "exact") + <> ( if planAgreesWithQuery q + then "" + else + " -- !! the captured plan produced " + <> Text.pack (show q.planTopRows) + <> " rows but the query returned " + <> Text.pack (show q.rowsReturned) + <> "; the EXPLAIN is describing a different query and cannot be trusted" + ), + q.planText + ] + +-- | Whether a captured plan used the HNSW index — the approximate path — as opposed to the +-- exact plan over @kioku_memories_scope_idx@. Which one Postgres picked is the load-bearing +-- observation: the previous initiative found it returning 50 correct rows on the exact plan and +-- zero on the HNSW one, and a measurement that does not record the plan has not measured the +-- thing that matters. +usedHnswIndex :: RecallQuality -> Bool +usedHnswIndex q = "kioku_memories_embedding_hnsw" `isInfixOf` Text.unpack q.planText + +-- | Run arbitrary DDL against the seeded corpus and re-@ANALYZE@ — the seam the bake-off needs +-- to prototype an /index/ candidate without shipping a migration for it. Each statement is its +-- own committed transaction, because a rebuilt index that the planner cannot see is the same +-- confidently-wrong measurement as an uncommitted row. +runDdl :: (Store :> es) => [Text] -> Eff es () +runDdl statements = do + traverse_ (runTransaction . Tx.sql . encodeUtf8) statements + runTransaction (Tx.sql "ANALYZE kioku_memories") diff --git a/test/Kioku/RecallSpec.hs b/test/Kioku/RecallSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/RecallSpec.hs @@ -0,0 +1,123 @@ +module Kioku.RecallSpec + ( tests, + ) +where + +import Data.Set qualified as Set +import Data.Text (Text) +import Data.Time (UTCTime, addUTCTime) +import Kioku.Api.Scope (MemoryScope (..), Namespace (..)) +import Kioku.Api.Types (MemoryRecord (..)) +import Kioku.Recall +import Kioku.Recall.Capability (CapabilityProbe (..), VectorCapability (..), classifyProbe) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Recall scoring" + [ testCase "RRF fusion favors a memory present in both lists" testRrfFusion, + testCase "signal blending maps recency priority and confidence" testSignalBlending, + testCase "character budgets truncate and stop before total cap" testBudgets, + testCase "execution plan fails open to keyword when vectors are unavailable" testFailOpenPlan, + testCase "capability classification names the reason vectors are unusable" testCapabilityClassification + ] + +testRrfFusion :: Assertion +testRrfFusion = do + let now = read "2026-06-24 00:00:00 UTC" + a = row "a" now "alpha" + b = row "b" now "bravo" + c = row "c" now "charlie" + hits = fuseRecallCandidates now [a, b] [b, c] + (hitId <$> hits) @?= ["b", "a", "c"] + assertApprox "rank 1 RRF term" (1 / 61) (rrfTerm 1) + +testSignalBlending :: Assertion +testSignalBlending = do + let now = read "2026-06-24 00:00:00 UTC" + thirtyDaysAgo = addUTCTime (negate (30 * 86400)) now + memoryRow = (row "m" thirtyDaysAgo "content") {priority = 0, confidence = "high"} + expected = rrfTerm 1 + rrfTerm 2 + (0.10 * 0.5) + (0.15 * 1.0) + (0.05 * 1.0) + assertApprox "recency halves at 30 days" 0.5 (recencyDecay now thirtyDaysAgo) + priorityWeight 0 @?= 1.0 + confidenceWeight "medium" @?= 0.6 + assertApprox "blended score" expected (blendScore now memoryRow (Just 1) (Just 2)) + +testBudgets :: Assertion +testBudgets = do + let now = read "2026-06-24 00:00:00 UTC" + hit1 = RecallHit {memory = row "a" now "abcdefghij", score = 2, ftsRank = Just 1, vecRank = Nothing} + hit2 = RecallHit {memory = row "b" now "klmnopqrst", score = 1, ftsRank = Just 2, vecRank = Nothing} + budgeted = applyCharacterBudgets 5 8 [hit1, hit2] + case budgeted of + [onlyHit] -> do + hitId onlyHit @?= "a" + onlyHit.memory.content @?= "ab..." + other -> fail ("Expected one budgeted hit, got " <> show (hitId <$> other)) + +testFailOpenPlan :: Assertion +testFailOpenPlan = do + planRecallExecution VectorAvailable Keyword + @?= RecallExecutionPlan {runFts = True, runVector = False, needsQueryEmbedding = False} + planRecallExecution VectorAvailable Embedding + @?= RecallExecutionPlan {runFts = False, runVector = True, needsQueryEmbedding = True} + planRecallExecution VectorAvailable Hybrid + @?= RecallExecutionPlan {runFts = True, runVector = True, needsQueryEmbedding = True} + planRecallExecution VectorExtensionUnavailable Hybrid + @?= RecallExecutionPlan {runFts = True, runVector = False, needsQueryEmbedding = False} + planRecallExecution (VectorColumnsUnavailable ["embedding"]) Embedding + @?= RecallExecutionPlan {runFts = True, runVector = False, needsQueryEmbedding = False} + -- A mismatch is a configuration error, not a missing feature, but recall's response is the + -- same: the vector channel cannot work, so degrade rather than fail on every query. + planRecallExecution (VectorDimensionMismatch 512 1536) Hybrid + @?= RecallExecutionPlan {runFts = True, runVector = False, needsQueryEmbedding = False} + +testCapabilityClassification :: Assertion +testCapabilityClassification = do + classifyProbe 1536 healthyProbe @?= VectorAvailable + -- The type must be nameable on *this* connection's search_path, not merely installed + -- somewhere in the database: recall casts with a bare `$1::vector`. + classifyProbe 1536 healthyProbe {hasVectorType = False} @?= VectorExtensionUnavailable + classifyProbe 1536 healthyProbe {hasEmbedding = False, embeddingTypmod = Nothing} + @?= VectorColumnsUnavailable ["embedding"] + classifyProbe 512 healthyProbe @?= VectorDimensionMismatch 512 1536 + -- A column declared without a width constrains nothing, so there is nothing to disagree + -- with; -1 is what pgvector reports for `vector` with no dimension. + classifyProbe 512 healthyProbe {embeddingTypmod = Just (-1)} @?= VectorAvailable + where + healthyProbe = + CapabilityProbe + { hasVectorType = True, + hasEmbedding = True, + hasEmbeddingModel = True, + hasDimensions = True, + hasContentHash = True, + embeddingTypmod = Just 1536 + } + +row :: Text -> UTCTime -> Text -> MemoryRecord +row memoryId createdAt content = + MemoryRecord + { memoryId, + agentId = "agent", + sessionId = Nothing, + scope = ScopeGlobal (Namespace "rei"), + memoryType = "pattern", + content, + priority = 100, + confidence = "medium", + tags = Set.empty, + status = "active", + createdAt + } + +hitId :: RecallHit -> Text +hitId hit = hit.memory.memoryId + +assertApprox :: String -> Double -> Double -> Assertion +assertApprox label expected actual = + assertBool + (label <> ": expected " <> show expected <> ", got " <> show actual) + (abs (expected - actual) < 0.0001) diff --git a/test/Kioku/RecallSqlSpec.hs b/test/Kioku/RecallSqlSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/RecallSqlSpec.hs @@ -0,0 +1,507 @@ +-- | Database-level tests for the recall candidate SQL: the scope predicate, the status +-- filter, the full-text query parser, and the vector round-trip. These go through +-- 'selectFtsCandidates' / 'selectVectorCandidates' rather than 'recall', so no embedding +-- endpoint is involved and the SQL under test is exercised exactly as it ships. +module Kioku.RecallSqlSpec (tests) where + +import Data.List (sort) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding (encodeUtf8) +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import Effectful (Eff, (:>)) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), ScopeKind (..)) +import Kioku.Api.Types (MemoryRecord (..)) +import Kioku.App (AppEffects, runAppIO, withNoopAppEnv) +import Kioku.Migrations.TestSupport (withKiokuMigratedDatabase) +import Kioku.Recall (RecallRequest (..), RecallStrategy (..), selectFtsCandidates, selectVectorCandidates, vectorLiteral) +import Kioku.Recall.Capability (VectorCapability (..), detectVectorCapability) +import Kioku.RecallHarness + ( CorpusConfig (..), + RecallQuality (..), + SeededCorpus (..), + cosineDistanceAtAngle, + defaultStarvationCorpus, + describeRecallQuality, + measureRecallQuality, + planAgreesWithQuery, + queryVector, + seedCorpus, + ) +import Kiroku.Store.Connection (defaultConnectionSettings) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Transaction (runTransaction) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertEqual, assertFailure, testCase) + +tests :: TestTree +tests = + testGroup + "Recall.Sql" + [ testCase "a global scope searches the whole namespace; an entity scope is exact" testScopePredicate, + testCase "archived memories are never candidates" testStatusFilter, + testGroup + "websearch_to_tsquery survives whatever the user typed" + [ testCase "an empty query" (assertQueryDoesNotThrow ""), + testCase "unbalanced quotes and bare operators" (assertQueryDoesNotThrow "\"unbalanced OR AND"), + testCase "punctuation only" (assertQueryDoesNotThrow "-- ; ()") + ], + testCase "a vector round-trip ranks the nearest embedding first" testVectorRoundTrip, + testCase "capability detection reads the column's real width" testDimensionDetection, + testCase "the harness seeds the geometry it claims" testHarnessGeometry, + testCase "the captured plan describes the query that was measured" testPlanCaptureIsFaithful, + testCase "the vector channel does not starve on a selective scope" testVectorChannelDoesNotStarve, + testCase "a healthy scope never pays for the exact fallback" testHealthyScopeSkipsFallback + ] + +-- | The defect, stated as the behaviour we want. +-- +-- The HNSW index is post-filtered: it picks candidates by distance alone, and the namespace, +-- scope, and status predicates are applied afterwards, to rows it has already chosen. When the +-- rows outside the scope are nearer the query than the in-scope answers — a small scope inside a +-- large namespace, which is the normal shape of kioku data — the index can spend its entire +-- budget on rows the filter then discards, and the vector channel returns nothing at all. +-- +-- Nothing notices. Recall fuses its two channels by rank, so a vector channel that returns zero +-- rows contributes zero ranks and the blended score degrades smoothly into pure keyword scoring: +-- no error, no warning, and nothing in a @RecallHit@ recording that the semantic half of a +-- "hybrid" search came back empty. The caller gets plausible keyword results and cannot tell. +-- That silence is why this defect survived, and it is why this case exists. +testVectorChannelDoesNotStarve :: IO () +testVectorChannelDoesNotStarve = + withStarvationCorpus \case + Nothing -> + putStrLn " [skipped] no reachable pgvector on this cluster; re-enter the dev shell to exercise the vector path" + Just q -> do + assertBool + ("the vector channel starved.\n" <> describeRecallQuality q) + (q.rowsReturned > 0) + assertBool + ("the vector channel returned candidates, but missed most of the true nearest.\n" <> describeRecallQuality q) + (q.recallAtK >= 0.5) + + -- The mechanism, pinned. Recall survives this corpus because the approximate pass starves + -- and the exact pass rescues it — not because the planner happened to choose kindly. If a + -- future change makes the approximate pass succeed here, that is good news, but it means + -- this case has stopped exercising the fallback it exists to prove, and the corpus needs to + -- be made harsher rather than the assertion relaxed. + assertBool + ( "the approximate pass did NOT starve on the starving corpus (" + <> show q.annRows + <> " rows), so this case is no longer exercising the fallback it was built to prove.\n" + <> describeRecallQuality q + ) + (q.annRows < 50) + assertBool + "the exact fallback did not fire, so recall survived this corpus by luck rather than by mechanism" + q.exactFallbackFired + +-- | Bar (b) of the bake-off, as a test: a fix that trades a rare failure for a common one is not +-- a fix. On a corpus with nothing out of scope there is nothing for the filter to discard, the +-- approximate pass fills its pool, and the exact fallback must not run — because the fallback +-- scans every embedded row in the caller's scope, and paying that on every healthy recall would +-- be a worse bug than the one being fixed. +-- +-- This also pins the @hnsw.ef_search@ setting. At pgvector's default (40, which is below the pool +-- size of 50) the approximate pass returns only 40 rows even here, the fallback would fire on +-- every single recall, and this case fails. +testHealthyScopeSkipsFallback :: IO () +testHealthyScopeSkipsFallback = + withRecallFixture \runEff -> do + result <- runEff do + available <- vectorTypeIsReachable + if not available + then pure Nothing + else do + corpus <- seedCorpus defaultStarvationCorpus {inScopeCount = 2000, decoyCount = 0} + Just <$> measureRecallQuality corpus 10 + case result of + Left err -> assertFailure ("store error: " <> show err) + Right Nothing -> + putStrLn " [skipped] no reachable pgvector on this cluster; re-enter the dev shell to exercise the vector path" + Right (Just q) -> do + assertEqual + "the approximate pass did not fill the candidate pool on a corpus with nothing to filter out" + 50 + q.annRows + assertBool + ( "the exact fallback fired on a healthy scope, which would make every recall pay for a \ + \full scan of its scope.\n" + <> describeRecallQuality q + ) + (not q.exactFallbackFired) + assertEqual "a healthy scope still returns a full pool" 50 q.rowsReturned + +-- | The instrument's self-check, on the corpus that actually matters, plus the one invariant +-- that must hold in every regime. +-- +-- This is /not/ marked expected-to-fail, and that is deliberate. It asserts what must be true +-- both before and after the starvation is fixed, so a regression in either property fails the +-- suite immediately rather than hiding inside the expected failure next door. +-- +-- The plan-agreement check earns its place. An @EXPLAIN@ whose SQL differs from the real +-- statement in any way the planner cares about can choose a different plan and report a +-- different answer — and the failure is silent and flattering. An earlier version of the harness +-- selected @memory_id@ alone instead of recall's thirteen columns; the narrower row made the +-- top-N sort look cheap, the planner took the exact plan, and the EXPLAIN cheerfully reported 50 +-- rows while the real query took the HNSW plan and returned zero. A measurement that disagrees +-- with the thing it claims to measure is worse than no measurement, so it is now an assertion. +testPlanCaptureIsFaithful :: IO () +testPlanCaptureIsFaithful = + withStarvationCorpus \case + Nothing -> + putStrLn " [skipped] no reachable pgvector on this cluster; re-enter the dev shell to exercise the vector path" + Just q -> do + assertBool + ( "the captured plan produced " + <> show q.planTopRows + <> " rows but the query returned " + <> show q.rowsReturned + <> "; the EXPLAIN is describing a query nobody runs and every number built on it is void.\n" + <> Text.unpack q.planText + ) + (planAgreesWithQuery q) + assertEqual + "an out-of-scope memory reached the caller: the scope filter is a correctness boundary, not a ranking hint" + 0 + q.decoysReturned + +-- | Seed 'defaultStarvationCorpus' and measure the vector channel against it, or report that +-- the cluster has no pgvector. A silent skip is indistinguishable from a pass, and this module +-- exists to make an invisible failure visible, so its own skips say so out loud. +withStarvationCorpus :: (Maybe RecallQuality -> IO ()) -> IO () +withStarvationCorpus assertOn = + withRecallFixture \runEff -> do + result <- runEff do + available <- vectorTypeIsReachable + if not available + then pure Nothing + else do + corpus <- seedCorpus defaultStarvationCorpus + Just <$> measureRecallQuality corpus 10 + case result of + Left err -> assertFailure ("store error: " <> show err) + Right measured -> assertOn measured + +-- | The instrument's own test, and the one that everything downstream rests on. +-- +-- 'Kioku.RecallHarness' claims that a row seeded at angle @t@ sits at cosine distance +-- @1 - cos t@ from the query, and it uses that claim to know the true ranking of the corpus +-- /without asking the database/ — which is what makes recall@k a measurement rather than a +-- tautology. If the claim is false, every number the harness reports is fiction, so it is +-- checked directly here against the distances Postgres actually computes. +-- +-- The second assertion is the one that makes starvation possible at all: every decoy must be +-- strictly nearer the query than every in-scope row, because starvation requires the filter to +-- correlate with distance. A corpus where that does not hold cannot starve, and a test built on +-- it would pass for the wrong reason. +testHarnessGeometry :: IO () +testHarnessGeometry = + withRecallFixture \runEff -> do + result <- runEff do + available <- vectorTypeIsReachable + if not available + then pure Nothing + else do + corpus <- seedCorpus smallCorpus + measured <- actualDistances + pure (Just (corpus, measured)) + case result of + Left err -> assertFailure ("store error: " <> show err) + Right Nothing -> + putStrLn " [skipped] no reachable pgvector on this cluster; re-enter the dev shell to exercise the recall harness" + Right (Just (corpus, measured)) -> do + assertEqual + "every seeded row came back" + (smallCorpus.inScopeCount + smallCorpus.decoyCount) + (length measured) + + -- 1. Postgres agrees with the closed form the ground truth is built on. + let predicted memoryId + | Just t <- lookup memoryId (seedAngles smallCorpus) = cosineDistanceAtAngle t + | otherwise = error ("unseeded memory id came back: " <> Text.unpack memoryId) + worstError = + maximum [abs (d - predicted memoryId) | (memoryId, d) <- measured] + assertBool + ( "Postgres disagrees with the harness's closed-form distance by " + <> show worstError + <> "; the ground truth cannot be trusted" + ) + -- pgvector stores each component as a 4-byte float, so agreement is to float + -- precision, not double. + (worstError < 1e-5) + + -- 2. Every decoy really is nearer than every in-scope row, or nothing can starve. + let distancesFor prefix = + [d | (memoryId, d) <- measured, prefix `Text.isPrefixOf` memoryId] + decoyDistances = distancesFor "m_decoy_" + inScopeDistances = distancesFor "m_in_" + assertBool + ( "the decoys are not strictly nearer than the in-scope rows: farthest decoy " + <> show (maximum decoyDistances) + <> " vs nearest in-scope " + <> show (minimum inScopeDistances) + <> " — this corpus cannot starve" + ) + (maximum decoyDistances < minimum inScopeDistances) + + -- 3. The ground-truth ordering is what it says it is: nearest first. + let truthDistances = + [d | memoryId <- corpus.trueNearestInScope, Just d <- [lookup memoryId measured]] + assertEqual + "the ground truth names every in-scope row" + smallCorpus.inScopeCount + (length truthDistances) + assertBool + "the ground truth is not ordered nearest-first" + (and (zipWith (<=) truthDistances (drop 1 truthDistances))) + +-- | Small enough to be quick, large enough that both bands are populated. The starvation case +-- uses 'defaultStarvationCorpus'; this one only has to prove the geometry. +smallCorpus :: CorpusConfig +smallCorpus = defaultStarvationCorpus {inScopeCount = 20, decoyCount = 20} + +-- | The angles the harness seeded, recomputed here from the same config, so the assertion +-- compares Postgres against the /closed form/ rather than against the harness's own vectors. +seedAngles :: CorpusConfig -> [(Text, Double)] +seedAngles cfg = + [ ("m_in_" <> Text.pack (show i), t) + | (i, t) <- zip [0 :: Int ..] (evenly cfg.inScopeCount cfg.inScopeAngles) + ] + <> [ ("m_decoy_" <> Text.pack (show i), t) + | (i, t) <- zip [0 :: Int ..] (evenly cfg.decoyCount cfg.decoyAngles) + ] + where + evenly n (lo, hi) + | n <= 0 = [] + | n == 1 = [lo] + | otherwise = [lo + (hi - lo) * fromIntegral i / fromIntegral (n - 1) | i <- [0 .. n - 1]] + +-- | Every seeded row's true cosine distance to the query, as Postgres computes it. The harness +-- must never do this — its whole value is knowing the answer without asking — but the test that +-- validates the harness must. +actualDistances :: (Store :> es) => Eff es [(Text, Double)] +actualDistances = + runTransaction (Tx.statement (vectorLiteral queryVector) stmt) + where + stmt :: Statement Text [(Text, Double)] + stmt = + preparable + "SELECT memory_id, (embedding <=> $1::vector)::float8 \ + \ FROM kiroku.kioku_memories \ + \ WHERE embedding IS NOT NULL \ + \ ORDER BY 2" + (E.param (E.nonNullable E.text)) + ( D.rowList + ( (,) + <$> D.column (D.nonNullable D.text) + <*> D.column (D.nonNullable D.float8) + ) + ) + +-- | The migrated schema declares @embedding vector(1536)@, so a process configured for 512 +-- is misconfigured and every embedding write it attempts would fail on the @::vector@ cast. +-- Detection reports it once, at startup, instead of once per event forever. +testDimensionDetection :: IO () +testDimensionDetection = + withRecallFixture \runEff -> do + result <- runEff do + available <- vectorTypeIsReachable + if not available + then pure Nothing + else do + mismatched <- detectVectorCapability 512 + matched <- detectVectorCapability 1536 + pure (Just (mismatched, matched)) + case result of + Left err -> assertFailure ("store error: " <> show err) + Right Nothing -> + putStrLn " [skipped] no reachable pgvector on this cluster; re-enter the dev shell to exercise dimension detection" + Right (Just (mismatched, matched)) -> do + assertEqual "a 512-dimension config against a vector(1536) column is a mismatch" (VectorDimensionMismatch 512 1536) mismatched + assertEqual "the configured width matching the column is simply available" VectorAvailable matched + +-- | Recall's scope predicate is @(($3 IS NULL AND $4 IS NULL) OR (scope_kind = $3 AND +-- scope_ref = $4))@: for a global scope both parameters are NULL, the first disjunct is +-- always true, and the query searches the entire namespace -- entity-scoped rows included. +-- That is deliberate and is *not* what the scoped read-model queries do with the same +-- 'MemoryScope' value. See docs/user/recall.md. +testScopePredicate :: IO () +testScopePredicate = + withRecallFixture \runEff -> do + result <- runEff do + seedMemories + [ ("m_global", ns1Global, "the deployment pipeline runs on nix flakes", "active"), + ("m_entity", ns1Entity, "the deployment pipeline runs on nix flakes", "active"), + ("m_other_ns", ns2Global, "the deployment pipeline runs on nix flakes", "active") + ] + wide <- selectFtsCandidates (request ns1Global "deployment pipeline") + exact <- selectFtsCandidates (request ns1Entity "deployment pipeline") + pure (memoryIds wide, memoryIds exact) + case result of + Left err -> assertFailure ("store error: " <> show err) + Right (wide, exact) -> do + assertEqual + "a global scope recalls every active row in the namespace, entity-scoped rows included" + ["m_entity", "m_global"] + wide + assertEqual + "an entity scope recalls only rows carrying exactly that scope" + ["m_entity"] + exact + assertBool + "no query ever crosses a namespace boundary" + ("m_other_ns" `notElem` (wide <> exact)) + +testStatusFilter :: IO () +testStatusFilter = + withRecallFixture \runEff -> do + result <- runEff do + seedMemories + [ ("m_active", ns1Global, "the deployment pipeline runs on nix flakes", "active"), + ("m_archived", ns1Global, "the deployment pipeline runs on nix flakes", "archived") + ] + memoryIds <$> selectFtsCandidates (request ns1Global "deployment pipeline") + case result of + Left err -> assertFailure ("store error: " <> show err) + Right ids -> assertEqual "an archived memory matching the query is not a candidate" ["m_active"] ids + +-- | @websearch_to_tsquery@ is total by design -- it never raises on malformed input, unlike +-- @to_tsquery@. This pins that, because recall passes user text straight into it. +assertQueryDoesNotThrow :: Text -> IO () +assertQueryDoesNotThrow query = + withRecallFixture \runEff -> do + result <- runEff do + seedMemories [("m_active", ns1Global, "the deployment pipeline runs on nix flakes", "active")] + selectFtsCandidates (request ns1Global query) + case result of + Left err -> assertFailure ("recall raised on query " <> show query <> ": " <> show err) + Right _ -> pure () + +-- | Writes a real @vector(1536)@ through 'vectorLiteral' and reads it back through the +-- shipping ORDER BY, so the text encoding and the @::vector@ cast are both exercised. +-- Skipped where the cluster has no reachable pgvector, which stays a supported degradation. +testVectorRoundTrip :: IO () +testVectorRoundTrip = + withRecallFixture \runEff -> do + result <- runEff do + available <- vectorTypeIsReachable + if not available + then pure Nothing + else do + seedMemories + [ ("m_near", ns1Global, "orbital mechanics", "active"), + ("m_far", ns1Global, "orbital mechanics", "active") + ] + setEmbedding "m_near" (unitVector 0) + setEmbedding "m_far" (unitVector 1) + Just . memoryIdsInOrder <$> selectVectorCandidates (request ns1Global "orbital") (unitVector 0) + case result of + Left err -> assertFailure ("store error: " <> show err) + Right Nothing -> + -- A silent skip is indistinguishable from a pass, so say so out loud. ephemeral-pg + -- takes its binaries from PATH: a shell entered before pgvector was added to + -- nix/haskell.nix still spins up clusters without it. + putStrLn " [skipped] no reachable pgvector on this cluster; re-enter the dev shell to exercise the vector path" + Right (Just ids) -> + assertEqual + "the embedding nearest the query vector ranks first" + ["m_near", "m_far"] + ids + +-- * Fixture + +ns1Global, ns1Entity, ns2Global :: MemoryScope +ns1Global = ScopeGlobal (Namespace "ns1") +ns1Entity = ScopeEntity (Namespace "ns1") (ScopeKind "repo") "web" +ns2Global = ScopeGlobal (Namespace "ns2") + +request :: MemoryScope -> Text -> RecallRequest +request scope query = + RecallRequest {scope, query, strategy = Keyword, maxResults = 10} + +memoryIds :: [MemoryRecord] -> [Text] +memoryIds = sort . fmap (\row -> row.memoryId) + +memoryIdsInOrder :: [MemoryRecord] -> [Text] +memoryIdsInOrder = fmap (\row -> row.memoryId) + +-- | A 1536-dimension basis vector. Two distinct basis vectors are cosine-orthogonal +-- (distance 1) while a vector is at distance 0 from itself, so ranking is unambiguous. +unitVector :: Int -> Vector Double +unitVector i = + Vector.generate 1536 (\j -> if j == i then 1 else 0) + +seedMemories :: + (Store :> es) => + [(Text, MemoryScope, Text, Text)] -> + Eff es () +seedMemories rows = + runTransaction . Tx.sql . encodeUtf8 $ + "INSERT INTO kioku_memories (memory_id, agent_id, namespace, scope_kind, scope_ref, memory_type, content, status, created_at, updated_at) VALUES " + <> Text.intercalate ", " (row <$> rows) + where + row (memoryId, scope, content, status) = + "('" + <> memoryId + <> "', 'agent', '" + <> namespaceOf scope + <> "', " + <> sqlText (kindOf scope) + <> ", " + <> sqlText (refOf scope) + <> ", 'fact', '" + <> content + <> "', '" + <> status + <> "', now(), now())" + + namespaceOf = \case + ScopeGlobal (Namespace ns) -> ns + ScopeEntity (Namespace ns) _ _ -> ns + kindOf = \case + ScopeGlobal _ -> Nothing + ScopeEntity _ (ScopeKind k) _ -> Just k + refOf = \case + ScopeGlobal _ -> Nothing + ScopeEntity _ _ r -> Just r + + sqlText Nothing = "NULL" + sqlText (Just value) = "'" <> value <> "'" + +setEmbedding :: (Store :> es) => Text -> Vector Double -> Eff es () +setEmbedding memoryId embedding = + runTransaction . Tx.sql . encodeUtf8 $ + "UPDATE kioku_memories SET embedding = '" + <> vectorLiteral embedding + <> "'::vector WHERE memory_id = '" + <> memoryId + <> "'" + +-- | The question that actually matters is not "is the extension installed somewhere" but +-- "can this connection name the type" -- which is what recall's @$1::vector@ cast needs, and +-- what @to_regtype@ answers against the live search_path. +vectorTypeIsReachable :: (Store :> es) => Eff es Bool +vectorTypeIsReachable = + runTransaction (Tx.statement () stmt) + where + stmt :: Statement () Bool + stmt = + preparable + "SELECT to_regtype('vector') IS NOT NULL" + E.noParams + (D.singleRow (D.column (D.nonNullable D.bool))) + +-- | A migrated throwaway database plus a store, handed to the case as a runner. Each case +-- gets its own cluster, so the seeds cannot collide. +withRecallFixture :: ((forall a. Eff AppEffects a -> IO (Either StoreError a)) -> IO ()) -> IO () +withRecallFixture use = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) \env -> + use (runAppIO env) diff --git a/test/Kioku/ReiCompatSpec.hs b/test/Kioku/ReiCompatSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/ReiCompatSpec.hs @@ -0,0 +1,214 @@ +module Kioku.ReiCompatSpec + ( tests, + ) +where + +import Data.Aeson (eitherDecode) +import Data.ByteString.Lazy qualified as LBS +import Data.Maybe (fromMaybe) +import Data.Set qualified as Set +import Data.Time (UTCTime) +import Data.Time.Format.ISO8601 (iso8601ParseM) +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), ScopeKind (..)) +import Kioku.Api.Types (Confidence (..), MemoryType (..)) +import Kioku.Id (idText) +import Kioku.Memory.Domain (MemoryEvent (..), MemoryRecordedData (..)) +import Kioku.Memory.EventStream (parseMemoryEvent) +import Kioku.Session.Domain + ( InteractiveSessionRecordedData (..), + SessionCompletedData (..), + SessionEvent (..), + SessionFailedData (..), + SessionResumedData (..), + SessionStartedData (..), + ) +import Kioku.Session.EventStream (parseSessionEvent) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Rei legacy codec compatibility" + [ testCase "decodes agent_memory_recorded" do + value <- either fail pure (eitherDecode reiMemoryRecordedJson) + event <- either (fail . show) pure (parseMemoryEvent value) + case event of + MemoryRecorded d -> assertMemoryRecorded d + other -> fail ("Expected MemoryRecorded, got " <> show other), + testCase "decodes agent_session_started" do + value <- either fail pure (eitherDecode reiSessionStartedJson) + event <- either (fail . show) pure (parseSessionEvent value) + case event of + SessionStarted d -> assertSessionStarted d + other -> fail ("Expected SessionStarted, got " <> show other), + testCase "decodes agent_session_completed" do + event <- decodeSession reiSessionCompletedJson + case event of + SessionCompleted d -> do + idText d.sessionId @?= "kioku_session_01kvxa7d2cezhs874g3n8dfgme" + d.completedAt @?= at "2026-06-24T21:30:00Z" + d.modelUsed @?= Just "claude-opus-4-8" + d.summary @?= Just "planned the day" + other -> fail ("Expected SessionCompleted, got " <> show other), + testCase "decodes agent_session_failed" do + event <- decodeSession reiSessionFailedJson + case event of + SessionFailed d -> do + idText d.sessionId @?= "kioku_session_01kvxa7d2cezhs874g3n8dfgme" + d.failedAt @?= at "2026-06-24T21:30:00Z" + d.errorMessage @?= "model timed out" + other -> fail ("Expected SessionFailed, got " <> show other), + testCase "decodes interactive_session_recorded" do + event <- decodeSession reiInteractiveSessionRecordedJson + case event of + InteractiveSessionRecorded d -> do + idText d.sessionId @?= "kioku_session_01kvxa7d2cezhs874g3n8dfgme" + d.agentId @?= "demo-agent" + d.focus @?= "general_coaching" + d.scope @?= ScopeEntity (Namespace "rei") (ScopeKind "intention") "intention_demo" + d.startedAt @?= at "2026-06-24T20:10:00Z" + other -> fail ("Expected InteractiveSessionRecorded, got " <> show other), + -- The two cases below pin the upcast rule for native SessionResumed events written + -- before the `force` field existed. Under the old code an omitted correlation key + -- bypassed matching entirely, so a keyless legacy resume must replay through the + -- force arm of the new guard and a keyed one through the matching arm. Get this + -- backwards and historical streams stop hydrating. + testCase "a pre-force keyless resume decodes as a force resume" do + event <- decodeSession (resumedWithoutForceJson "null") + case event of + SessionResumed d -> do + d.correlationKey @?= Nothing + d.force @?= True + other -> fail ("Expected SessionResumed, got " <> show other), + testCase "a pre-force keyed resume decodes as a plain resume" do + event <- decodeSession (resumedWithoutForceJson "\"k1\"") + case event of + SessionResumed d -> do + d.correlationKey @?= Just "k1" + d.force @?= False + other -> fail ("Expected SessionResumed, got " <> show other) + ] + +decodeSession :: LBS.ByteString -> IO SessionEvent +decodeSession raw = do + value <- either fail pure (eitherDecode raw) + either (fail . show) pure (parseSessionEvent value) + +at :: String -> UTCTime +at = fromMaybe (error "bad fixture timestamp") . iso8601ParseM + +assertMemoryRecorded :: MemoryRecordedData -> Assertion +assertMemoryRecorded d = do + idText d.memoryId @?= "kioku_memory_01kvx9my35e5y825cpy4nycjgz" + (idText <$> d.sessionId) @?= Just "kioku_session_01kvxa7d2cezhs874g3n8dfgme" + d.agentId @?= "demo-agent" + d.scope @?= ScopeEntity (Namespace "rei") (ScopeKind "intention") "intention_demo" + d.memoryType @?= MemoryPreference + d.content @?= "prefers concise answers" + d.priority @?= 100 + d.confidence @?= HighConfidence + d.tags @?= Set.fromList ["style"] + (idText <$> d.supersedes) @?= Just "kioku_memory_01kvx9pkrxevzafwat8k704yzh" + +assertSessionStarted :: SessionStartedData -> Assertion +assertSessionStarted d = do + idText d.sessionId @?= "kioku_session_01kvxa7d2cezhs874g3n8dfgme" + d.agentId @?= "demo-agent" + d.focus @?= "today" + d.scope @?= ScopeEntity (Namespace "rei") (ScopeKind "intention") "intention_demo" + d.subjectRef @?= Just "daily planning" + (idText <$> d.previousSessionId) @?= Just "kioku_session_01kvxa2gw6er7r2yzpvtq9axch" + +reiMemoryRecordedJson :: LBS.ByteString +reiMemoryRecordedJson = + """ + { + "type": "agent_memory_recorded", + "data": { + "memoryId": "agent_memory_01kvx9my35e5y825cpy4nycjgz", + "agentId": "demo-agent", + "sessionId": "agent_session_01kvxa7d2cezhs874g3n8dfgme", + "memoryType": "preference", + "content": "prefers concise answers", + "anchor": { + "type": "intention", + "id": "intention_demo" + }, + "confidence": "high", + "tags": ["style"], + "supersedes": "agent_memory_01kvx9pkrxevzafwat8k704yzh", + "recordedAt": "2026-06-24T20:10:00Z" + } + } + """ + +reiSessionStartedJson :: LBS.ByteString +reiSessionStartedJson = + """ + { + "type": "agent_session_started", + "data": { + "sessionId": "agent_session_01kvxa7d2cezhs874g3n8dfgme", + "agentId": "demo-agent", + "focusType": "FocusToday", + "intentionId": "intention_demo", + "previousSessionId": "agent_session_01kvxa2gw6er7r2yzpvtq9axch", + "focusTarget": "daily planning", + "startedAt": "2026-06-24T20:10:00Z" + } + } + """ + +reiSessionCompletedJson :: LBS.ByteString +reiSessionCompletedJson = + """ + { + "type": "agent_session_completed", + "data": { + "sessionId": "agent_session_01kvxa7d2cezhs874g3n8dfgme", + "completedAt": "2026-06-24T21:30:00Z", + "modelUsed": "claude-opus-4-8", + "summary": "planned the day" + } + } + """ + +reiSessionFailedJson :: LBS.ByteString +reiSessionFailedJson = + """ + { + "type": "agent_session_failed", + "data": { + "sessionId": "agent_session_01kvxa7d2cezhs874g3n8dfgme", + "failedAt": "2026-06-24T21:30:00Z", + "errorMessage": "model timed out" + } + } + """ + +reiInteractiveSessionRecordedJson :: LBS.ByteString +reiInteractiveSessionRecordedJson = + """ + { + "type": "interactive_session_recorded", + "data": { + "sessionId": "agent_session_01kvxa7d2cezhs874g3n8dfgme", + "agentId": "demo-agent", + "focusType": "FocusGeneralCoaching", + "intentionId": "intention_demo", + "startedAt": "2026-06-24T20:10:00Z" + } + } + """ + +-- | A native @SessionResumed@ payload as written before the @force@ field existed. +resumedWithoutForceJson :: LBS.ByteString -> LBS.ByteString +resumedWithoutForceJson correlationKey = + "{\"type\": \"session_resumed\", \"data\": {" + <> "\"sessionId\": \"kioku_session_01kvxa7d2cezhs874g3n8dfgme\", " + <> "\"correlationKey\": " + <> correlationKey + <> ", " + <> "\"input\": \"approved\", " + <> "\"resumedAt\": \"2026-06-24T21:30:00Z\"}}" diff --git a/test/Kioku/SchemaSpec.hs b/test/Kioku/SchemaSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/SchemaSpec.hs @@ -0,0 +1,131 @@ +-- | Database-level tests for the read-model schema itself: the constraints and indexes +-- the migrations are supposed to install. These go through a raw hasql connection rather +-- than the kiroku 'Store' effect, because the store's transaction error mapper folds a +-- unique violation on an application table into 'WrongExpectedVersion' (its 23505 branch +-- is written for the event tables) and so cannot tell us the SQLSTATE we care about. +module Kioku.SchemaSpec (tests) where + +import Control.Exception (bracket) +import Data.Text (Text) +import Data.Text qualified as Text +import Hasql.Connection qualified as Connection +import Hasql.Connection.Settings qualified as Settings +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Errors (ServerError (..), SessionError (..), StatementError (..)) +import Hasql.Session qualified as Session +import Hasql.Statement (Statement, preparable) +import Kioku.Migrations.TestSupport (withKiokuMigratedDatabase) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Schema" + [ testGroup + "global scopes are unique, not merely NULL-distinct" + [ testCase "two global scenes with the same key collide" (assertSqlState "23505" duplicateGlobalScenes), + testCase "two global personas in one namespace collide" (assertSqlState "23505" duplicateGlobalPersonas) + ], + testGroup + "a scope is global or an entity scope, never half of one" + [ testCase "a memory with a kind and no ref is rejected" (assertSqlState "23514" (halfScopedMemory "'kind-without-ref'" "NULL")), + testCase "a memory with a ref and no kind is rejected" (assertSqlState "23514" (halfScopedMemory "NULL" "'ref-without-kind'")) + ], + testCase "the chain and session-list indexes exist and the redundant turns index is gone" testIndexes + ] + +-- * Constraint cases + +-- | Both rows carry distinct primary keys, so only the scope constraint can reject the +-- second one. Under the old @UNIQUE (namespace, scope_kind, scope_ref, scene_key)@ these +-- two rows coexisted happily: SQL considers NULLs distinct, so the constraint enforced +-- nothing at all for global scopes. +duplicateGlobalScenes :: Text +duplicateGlobalScenes = + """ + INSERT INTO kiroku.kioku_scenes (scene_id, namespace, scene_key, title, body_md, source_hash) + VALUES ('scene-a', 'ns', 'default', 't', 'b', 'h'), + ('scene-b', 'ns', 'default', 't', 'b', 'h') + """ + +duplicateGlobalPersonas :: Text +duplicateGlobalPersonas = + """ + INSERT INTO kiroku.kioku_personas (persona_id, namespace, body_md, source_hash) + VALUES ('persona-a', 'ns', 'b', 'h'), + ('persona-b', 'ns', 'b', 'h') + """ + +-- | A row with exactly one of the two scope columns set. 'Kioku.Api.Scope.scopeFromColumns' +-- reads it back as a global scope, yet no exact-scope query matches it -- so the row is +-- invisible to both halves of the API. The CHECK makes it unwritable. Both arguments are +-- SQL literals, so a case can pass @NULL@ for the column it wants to leave unset. +halfScopedMemory :: Text -> Text -> Text +halfScopedMemory scopeKind scopeRef = + "INSERT INTO kiroku.kioku_memories" + <> " (memory_id, agent_id, namespace, scope_kind, scope_ref, memory_type, content, created_at, updated_at)" + <> " VALUES ('m-half', 'agent', 'ns', " + <> scopeKind + <> ", " + <> scopeRef + <> ", 'fact', 'content', now(), now())" + +assertSqlState :: Text -> Text -> IO () +assertSqlState expected sql = + withMigratedConnection \conn -> do + result <- Connection.use conn (Session.script sql) + case result of + Right () -> + assertFailure + ("expected SQLSTATE " <> Text.unpack expected <> ", but the statement succeeded: " <> Text.unpack sql) + Left err -> + case sqlState err of + Just actual -> actual @?= expected + Nothing -> assertFailure ("expected SQLSTATE " <> Text.unpack expected <> ", got: " <> show err) + +sqlState :: SessionError -> Maybe Text +sqlState = \case + ScriptSessionError _ (ServerError code _ _ _ _) -> Just code + StatementSessionError _ _ _ _ _ (ServerStatementError (ServerError code _ _ _ _)) -> Just code + _ -> Nothing + +-- * Index case + +testIndexes :: IO () +testIndexes = + withMigratedConnection \conn -> do + result <- Connection.use conn (Session.statement () selectKiokuIndexes) + case result of + Left err -> assertFailure ("listing indexes failed: " <> show err) + Right indexes -> do + mapM_ + (\name -> assertBool (Text.unpack name <> " is missing") (name `elem` indexes)) + [ "kioku_memories_supersedes_idx", + "kioku_memories_superseded_by_idx", + "kioku_sessions_namespace_started_idx", + "kioku_sessions_namespace_focus_idx" + ] + assertBool + "kioku_turns_session_idx still exists; it duplicates the index implied by UNIQUE (session_id, turn_index)" + ("kioku_turns_session_idx" `notElem` indexes) + +-- | @pg_indexes.indexname@ is a @name@, not a @text@; the cast is what lets hasql decode it. +selectKiokuIndexes :: Statement () [Text] +selectKiokuIndexes = + preparable + "SELECT indexname::text FROM pg_indexes WHERE schemaname = 'kiroku'" + E.noParams + (D.rowList (D.column (D.nonNullable D.text))) + +-- * Harness + +withMigratedConnection :: (Connection.Connection -> IO a) -> IO a +withMigratedConnection use = + withKiokuMigratedDatabase \connStr -> + bracket (acquire connStr) Connection.release use + where + acquire connStr = + Connection.acquire (Settings.connectionString connStr) + >>= either (\err -> assertFailure ("could not connect: " <> show err)) pure diff --git a/test/Kioku/ScopeIdentitySpec.hs b/test/Kioku/ScopeIdentitySpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/ScopeIdentitySpec.hs @@ -0,0 +1,111 @@ +module Kioku.ScopeIdentitySpec (tests) where + +import Data.List (nub) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Time (UTCTime) +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), ScopeKind (..), mkNamespace, mkScopeKind) +import Kioku.Distill.L2 (l2SceneTimerId, sceneRowId) +import Kioku.Distill.L3 (l3PersonaTimerId, personaRowId) +import Kioku.Distill.ScopeIdentity (escapeScopeComponent, scopeIdentity, scopeSlugFromColumns) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Scope identity" + [ testCase "two scopes that used to collide now derive different everything" testCollision, + testCase "well-formed scopes keep their exact legacy ids" testLegacyStability, + testCase "escaping is injective on adversarial components" testEscapeInjective, + testCase "mirror slugs separate scopes the sanitiser cannot" testSlugCollision, + testCase "namespace and kind reject the reserved characters" testValidators + ] + +-- | The canonical collision. Both of these used to render @a/b/c@, so they shared one scene +-- row, one persona row, one timer id and one mirror file — and the upserts do not update the +-- scope columns on conflict, so the second scope's content landed on the first scope's row. +collidingGlobal, collidingEntity :: MemoryScope +collidingGlobal = ScopeGlobal (Namespace "a/b/c") +collidingEntity = ScopeEntity (Namespace "a") (ScopeKind "b") "c" + +testCollision :: Assertion +testCollision = do + scopeIdentity collidingGlobal @?= "a%2Fb%2Fc" + scopeIdentity collidingEntity @?= "a/b/c" + + assertAllDistinct "scene row id" [sceneRowId collidingGlobal, sceneRowId collidingEntity] + assertAllDistinct "persona row id" [personaRowId collidingGlobal, personaRowId collidingEntity] + assertAllDistinct + "scene timer id" + [show (l2SceneTimerId collidingGlobal "src"), show (l2SceneTimerId collidingEntity "src")] + assertAllDistinct + "persona timer id" + [show (l3PersonaTimerId collidingGlobal fireAt), show (l3PersonaTimerId collidingEntity fireAt)] + assertAllDistinct + "mirror slug" + [ scopeSlugFromColumns "a/b/c" Nothing Nothing, + scopeSlugFromColumns "a" (Just "b") (Just "c") + ] + where + fireAt :: UTCTime + fireAt = read "2026-07-11 00:00:00 UTC" + +-- | Every scope observed in the hosts and the docs contains none of @%@, @/@, @:@, and so +-- must encode to itself. This is what lets the id-recompute migration touch almost nothing: +-- if these bytes changed, every existing scene and persona row would be orphaned. +testLegacyStability :: Assertion +testLegacyStability = do + sceneRowId (ScopeEntity (Namespace "rei") (ScopeKind "intention") "intention_abc") + @?= "kioku_scene:rei/intention/intention_abc:default" + personaRowId (ScopeEntity (Namespace "mori") (ScopeKind "repo") "web") + @?= "kioku_persona:mori/repo/web" + personaRowId (ScopeGlobal (Namespace "shikigami")) + @?= "kioku_persona:shikigami" + +-- | Refs are host free text and legitimately contain @/@ (repo-style refs such as +-- @shinzui/kikan@), so the encoding has to stay injective rather than reject them. +testEscapeInjective :: Assertion +testEscapeInjective = do + escapeScopeComponent "a/b" @?= "a%2Fb" + escapeScopeComponent "a%2Fb" @?= "a%252Fb" + escapeScopeComponent "a:b" @?= "a%3Ab" + escapeScopeComponent "plain-ref_123" @?= "plain-ref_123" + assertAllDistinct + "escaped adversarial components" + (escapeScopeComponent <$> ["a/b", "a%2Fb", "a:b", "a%3Ab", "a%25b"]) + +-- | The readable half of a slug cannot be collision-free: the sanitiser maps every unsafe +-- character to @-@, so these two scopes both render @a-b@. The hash suffix is what separates +-- them. +testSlugCollision :: Assertion +testSlugCollision = do + let dashedNamespace = scopeSlugFromColumns "a-b" Nothing Nothing + splitScope = scopeSlugFromColumns "a" (Just "b") Nothing + global = scopeSlugFromColumns "a/b/c" Nothing Nothing + entity = scopeSlugFromColumns "a" (Just "b") (Just "c") + assertAllDistinct "slugs whose readable prefixes both sanitise to a-b" [dashedNamespace, splitScope] + assertAllDistinct "slugs for the scopes that used to collide" [global, entity] + assertBool + ("the slug keeps a human-readable prefix, got " <> show entity) + ("a-b-c-" `Text.isPrefixOf` entity) + +testValidators :: Assertion +testValidators = do + mkNamespace "rei" @?= Right (Namespace "rei") + mkScopeKind "intention" @?= Right (ScopeKind "intention") + assertLeft "an empty namespace" (mkNamespace "") + assertLeft "a namespace with a slash" (mkNamespace "a/b") + assertLeft "a namespace with a percent" (mkNamespace "a%b") + assertLeft "a kind with a colon" (mkScopeKind "a:b") + +assertLeft :: (Show a) => String -> Either Text a -> Assertion +assertLeft label = \case + Left _ -> pure () + Right value -> assertBool (label <> " should have been rejected, got " <> show value) False + +assertAllDistinct :: (Eq a, Show a) => String -> [a] -> Assertion +assertAllDistinct label values = + assertBool + (label <> ": expected all distinct, got " <> show values) + (length (nub values) == length values) diff --git a/test/Kioku/SessionInvariantsSpec.hs b/test/Kioku/SessionInvariantsSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/SessionInvariantsSpec.hs @@ -0,0 +1,530 @@ +-- | Invariants the session aggregate enforces itself, independent of any read-model +-- precheck. Every test here that matters drives 'runCommandWithProjections' directly, so a +-- passing assertion proves the state machine rejected (or accepted) the command — not that +-- a precheck happened to run first. +module Kioku.SessionInvariantsSpec (tests) where + +import Control.Monad (void) +import Data.Aeson (Value, object, toJSON, (.=)) +import Data.Text (Text) +import Data.Time (UTCTime (..), fromGregorian, getCurrentTime) +import Data.Vector qualified as Vector +import Effectful (Eff, IOE, liftIO, (:>)) +import Effectful.Error.Static (Error) +import Keiro.Command (CommandError (..), defaultRunCommandOptions) +import Keiro.Projection (runCommandWithProjections) +import Keiro.Stream qualified as Stream +import Kioku.Api.Scope (MemoryScope (..), Namespace (..)) +import Kioku.App (AppEffects, runAppIO, withNoopAppEnv) +import Kioku.Distill.Timer (l1TimerScheduleProjection) +import Kioku.Id (SessionId, genSessionId, idText) +import Kioku.Migrations.TestSupport (withKiokuMigratedDatabase) +import Kioku.Session qualified as Session +import Kioku.Session.Domain + ( AwaitInputData (..), + RecordTurnData (..), + ResumeSessionData (..), + SessionAwaitingData (..), + SessionCommand (..), + SessionEvent (..), + SessionResumedData (..), + SessionStartedData (..), + StartSessionData (..), + ) +import Kioku.Session.EventStream (parseSessionEvent, sessionEventStream, sessionStream) +import Kioku.Session.ReadModel (SessionRow (..), sessionInlineProjection) +import Kiroku.Store.Append (appendToStream) +import Kiroku.Store.Connection (defaultConnectionSettings) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Read (readStreamForward) +import Kiroku.Store.Types + ( EventData (..), + EventType (..), + ExpectedVersion (..), + RecordedEvent (..), + StreamVersion (..), + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, assertFailure, testCase) + +tests :: TestTree +tests = + testGroup + "Session invariants" + [ testCase "aggregate rejects a mismatched resume key" testAggregateRejectsMismatchedKey, + testCase "a stale resume after re-park is rejected" testStaleResumeAfterRepark, + testCase "forceResume waives the key explicitly" testForceResume, + testCase "a keyless wait resumes with no key" testKeylessWaitResumes, + testCase "a keyed resume of a keyless wait is rejected" testKeyedResumeOfKeylessWait, + testCase "re-park clears the previous resume input" testReparkClearsResumeInput, + testCase "pre-force keyless resume events still hydrate" testHydratesLegacyKeylessResume, + testCase "pre-force keyed resume events still hydrate" testHydratesLegacyKeyedResume, + testGroup + "turn identity" + [ testCase "an identical re-record is a duplicate" testTurnReRecordIsDuplicate, + testCase "the same index with different content is a conflict" testTurnSameIndexConflict, + testCase "the same turn id at a different index is a conflict" testTurnIdReuseConflict, + testCase "the aggregate rejects a non-increasing index" testAggregateRejectsNonIncreasingTurn + ] + ] + +-- | The core proof: with the session parked on @k1@, a @ResumeSession@ carrying @k2@ is +-- rejected by the transducer itself. No read-model precheck runs on this path, so the +-- rejection can only have come from the aggregate's guard. +testAggregateRejectsMismatchedKey :: Assertion +testAggregateRejectsMismatchedKey = + withApp do + sid <- startFixture + parkFixture sid (Just "k1") + now <- liftIO getCurrentTime + result <- + runSessionCommandDirect + sid + ( ResumeSession + ResumeSessionData + { sessionId = sid, + correlationKey = Just "k2", + force = False, + input = "approved", + resumedAt = now + } + ) + liftIO $ assertRejected "mismatched resume key" result + events <- readSessionEvents sid + liftIO $ + assertEqual + "no SessionResumed event was appended" + ["SessionStarted", "SessionAwaiting"] + (eventName <$> events) + +-- | The race the plan closes: park k1, resume k1, re-park on k2. A caller still holding +-- k1 must not be able to answer the k2 wait. +testStaleResumeAfterRepark :: Assertion +testStaleResumeAfterRepark = + withApp do + sid <- startFixture + parkFixture sid (Just "k1") + resumeFixture sid (Just "k1") "first answer" + parkFixture sid (Just "k2") + now <- liftIO getCurrentTime + let staleResume = + ResumeSessionData + { sessionId = sid, + correlationKey = Just "k1", + force = False, + input = "stale answer", + resumedAt = now + } + -- Through the public API the precheck catches it early ... + apiResult <- Session.resume staleResume + case apiResult of + Left Session.SessionCorrelationMismatch -> pure () + other -> liftIO (assertFailure ("expected SessionCorrelationMismatch, got " <> show other)) + -- ... and with the precheck bypassed, the aggregate catches it too. + directResult <- runSessionCommandDirect sid (ResumeSession staleResume) + liftIO $ assertRejected "stale resume key" directResult + -- The live wait still resumes. + resumeFixture sid (Just "k2") "second answer" + row <- getExisting sid + liftIO do + assertEqual "resumed with the live key" "running" row.status + assertEqual "the live answer won" (Just "second answer") row.resumeInput + +testForceResume :: Assertion +testForceResume = + withApp do + sid <- startFixture + parkFixture sid (Just "k1") + now <- liftIO getCurrentTime + result <- Session.forceResume sid "forced answer" now + void (liftEither "Session.forceResume" result) + row <- getExisting sid + liftIO do + assertEqual "force resume ran" "running" row.status + assertEqual "force resume input" (Just "forced answer") row.resumeInput + events <- readSessionEvents sid + liftIO $ + assertBool "the appended SessionResumed event carries force = True" $ + case [d | SessionResumed d <- events] of + [d] -> d.force && isNothing' d.correlationKey + _ -> False + where + isNothing' = \case + Nothing -> True + Just (_ :: Text) -> False + +testKeylessWaitResumes :: Assertion +testKeylessWaitResumes = + withApp do + sid <- startFixture + parkFixture sid Nothing + resumeFixture sid Nothing "approved" + row <- getExisting sid + liftIO do + assertEqual "keyless resume ran" "running" row.status + assertEqual "keyless resume input" (Just "approved") row.resumeInput + +-- | Exact 'Maybe' equality: naming a key when the wait has none is a mismatch, not a +-- harmless extra. +testKeyedResumeOfKeylessWait :: Assertion +testKeyedResumeOfKeylessWait = + withApp do + sid <- startFixture + parkFixture sid Nothing + now <- liftIO getCurrentTime + result <- + runSessionCommandDirect + sid + ( ResumeSession + ResumeSessionData + { sessionId = sid, + correlationKey = Just "unexpected", + force = False, + input = "approved", + resumedAt = now + } + ) + liftIO $ assertRejected "keyed resume of a keyless wait" result + +testReparkClearsResumeInput :: Assertion +testReparkClearsResumeInput = + withApp do + sid <- startFixture + parkFixture sid (Just "k1") + resumeFixture sid (Just "k1") "first answer" + parkFixture sid (Just "k2") + row <- getExisting sid + liftIO do + assertEqual "re-parked" "awaiting" row.status + assertEqual "the new wait's key is visible" (Just "k2") row.awaitingCorrelationKey + assertEqual "the previous wait's answer is gone" Nothing row.resumeInput + +-- | A stream written before the @force@ field existed, whose resume omitted the +-- correlation key. Under the old code that omission /was/ a bypass, so it must replay +-- through the force arm of the new guard. +testHydratesLegacyKeylessResume :: Assertion +testHydratesLegacyKeylessResume = hydrateLegacyStream Nothing + +-- | The same, but the historical resume named the key it parked on: it must replay through +-- the matching arm. +testHydratesLegacyKeyedResume :: Assertion +testHydratesLegacyKeyedResume = hydrateLegacyStream (Just "k1") + +-- | Hand-append a pre-@force@ event stream (no @force@ key in the @SessionResumed@ JSON), +-- then drive a real command against it. Success proves keiro rehydrated all three events +-- through the new guard: a rejected historical event would surface as +-- 'HydrationReplayFailed' instead. +-- +-- The command runs through 'runCommandWithProjections' rather than 'Session.awaitInput' +-- because a hand-appended stream has no read-model row, and the public API's precheck +-- would fail with 'SessionNotFound' before hydration was ever attempted. +hydrateLegacyStream :: Maybe Text -> Assertion +hydrateLegacyStream resumedKey = + withApp do + sid <- liftIO genSessionId + now <- liftIO getCurrentTime + let started = + SessionStarted + SessionStartedData + { sessionId = sid, + agentId = "test-agent", + focus = "legacy replay", + scope = testScope, + subjectRef = Nothing, + previousSessionId = Nothing, + parentSessionId = Nothing, + delegationDepth = 0, + startedAt = now + } + awaiting = + SessionAwaiting + SessionAwaitingData + { sessionId = sid, + reason = "approval", + correlationKey = Just "k1", + deadline = Nothing, + awaitedAt = now + } + void $ + appendToStream + (Stream.streamName (sessionStream sid)) + NoStream + [ rawEvent "SessionStarted" (toJSON started), + rawEvent "SessionAwaiting" (toJSON awaiting), + rawEvent "SessionResumed" (legacyResumedPayload sid resumedKey now) + ] + -- The stream now ends in Running. A fresh AwaitInput must hydrate it and be accepted. + result <- + runSessionCommandDirect + sid + ( AwaitInput + AwaitInputData + { sessionId = sid, + reason = "approval", + correlationKey = Just "k2", + deadline = Nothing, + awaitedAt = now + } + ) + case result of + Right _ -> pure () + Left err -> + liftIO (assertFailure ("a pre-force stream failed to hydrate: " <> show err)) + events <- readSessionEvents sid + liftIO $ + assertEqual + "the new event landed on top of the replayed history" + ["SessionStarted", "SessionAwaiting", "SessionResumed", "SessionAwaiting"] + (eventName <$> events) + +-- * Turn identity + +testTurnReRecordIsDuplicate :: Assertion +testTurnReRecordIsDuplicate = + withApp do + sid <- startFixture + let turn = turnData sid "turn-a" 0 "hello" + void (liftEither "first recordTurn" =<< Session.recordTurn turn) + void (liftEither "duplicate recordTurn" =<< Session.recordTurn turn) + events <- readSessionEvents sid + liftIO $ + assertEqual + "the duplicate appended no second TurnRecorded" + 1 + (length [() | TurnRecorded _ <- events]) + turns <- liftEither "getTurns" =<< Session.getTurns sid + liftIO $ assertEqual "exactly one turn row" 1 (length turns) + +testTurnSameIndexConflict :: Assertion +testTurnSameIndexConflict = + withApp do + sid <- startFixture + void (liftEither "first recordTurn" =<< Session.recordTurn (turnData sid "turn-a" 0 "hello")) + result <- Session.recordTurn (turnData sid "turn-a" 0 "something else") + liftIO $ assertConflict "same index, different content" result + events <- readSessionEvents sid + liftIO $ + assertEqual "no second TurnRecorded" 1 (length [() | TurnRecorded _ <- events]) + +testTurnIdReuseConflict :: Assertion +testTurnIdReuseConflict = + withApp do + sid <- startFixture + void (liftEither "first recordTurn" =<< Session.recordTurn (turnData sid "turn-a" 0 "hello")) + result <- Session.recordTurn (turnData sid "turn-a" 1 "a new turn reusing the id") + liftIO $ assertConflict "turn id reused at a different index" result + events <- readSessionEvents sid + liftIO $ + assertEqual "no second TurnRecorded" 1 (length [() | TurnRecorded _ <- events]) + +-- | The command layer's dedup is a convenience; this proves the state machine enforces the +-- strictly-increasing index on its own, with every precheck bypassed. +testAggregateRejectsNonIncreasingTurn :: Assertion +testAggregateRejectsNonIncreasingTurn = + withApp do + sid <- startFixture + void (liftEither "recordTurn 0" =<< Session.recordTurn (turnData sid "turn-a" 0 "hello")) + void (liftEither "recordTurn 1" =<< Session.recordTurn (turnData sid "turn-b" 1 "again")) + now <- liftIO getCurrentTime + result <- + runSessionCommandDirect + sid + ( RecordTurn + RecordTurnData + { sessionId = sid, + turnId = "turn-c", + turnIndex = 1, + role = "user", + content = "a stale index", + toolSummary = Nothing, + promptTokens = Nothing, + outputTokens = Nothing, + recordedAt = now + } + ) + liftIO $ assertRejected "non-increasing turn index" result + events <- readSessionEvents sid + liftIO $ + assertEqual "still only the two committed turns" 2 (length [() | TurnRecorded _ <- events]) + +turnData :: SessionId -> Text -> Int -> Text -> RecordTurnData +turnData sid turnId turnIndex content = + RecordTurnData + { sessionId = sid, + turnId, + turnIndex, + role = "user", + content, + toolSummary = Nothing, + promptTokens = Nothing, + outputTokens = Nothing, + recordedAt = fixedRecordedAt + } + +-- | Turn recording compares semantic fields, not the clock (see 'Kioku.Session.mismatchOf'), +-- but pinning the timestamp keeps these fixtures obviously identical. +fixedRecordedAt :: UTCTime +fixedRecordedAt = UTCTime (fromGregorian 2026 7 11) 0 + +assertConflict :: String -> Either Session.SessionWriteError SessionId -> Assertion +assertConflict label = \case + Left (Session.SessionConflict _) -> pure () + other -> assertFailure (label <> ": expected SessionConflict, got " <> show other) + +-- | The @SessionResumed@ JSON exactly as it was written before this plan: no @force@ key. +legacyResumedPayload :: SessionId -> Maybe Text -> UTCTime -> Value +legacyResumedPayload sid key resumedAt = + object + [ "type" .= ("session_resumed" :: Text), + "data" + .= object + [ "sessionId" .= idText sid, + "correlationKey" .= key, + "input" .= ("approved" :: Text), + "resumedAt" .= resumedAt + ] + ] + +rawEvent :: Text -> Value -> EventData +rawEvent typ payload = + EventData + { eventId = Nothing, + eventType = EventType typ, + payload, + metadata = Nothing, + causationId = Nothing, + correlationId = Nothing + } + +-- | Run a command straight at the aggregate, skipping 'Kioku.Session''s read-model +-- prechecks. This is the harness that makes "the aggregate enforces it" a testable claim. +runSessionCommandDirect :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + SessionCommand -> + Eff es (Either CommandError ()) +runSessionCommandDirect sid cmd = + void + <$> runCommandWithProjections + defaultRunCommandOptions + sessionEventStream + (sessionStream sid) + cmd + [sessionInlineProjection, l1TimerScheduleProjection] + +assertRejected :: String -> Either CommandError () -> Assertion +assertRejected label = \case + Left CommandRejected -> pure () + Left err -> assertFailure (label <> ": expected CommandRejected, got " <> show err) + Right () -> assertFailure (label <> ": expected CommandRejected, but the command was accepted") + +withApp :: Eff AppEffects a -> IO a +withApp action = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) \env -> do + result <- runAppIO env action + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right value -> pure value + +startFixture :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + Eff es SessionId +startFixture = do + sid <- liftIO genSessionId + now <- liftIO getCurrentTime + result <- + Session.start + StartSessionData + { sessionId = sid, + agentId = "test-agent", + focus = "session invariants", + scope = testScope, + subjectRef = Nothing, + previousSessionId = Nothing, + parentSessionId = Nothing, + delegationDepth = 0, + startedAt = now + } + void (liftEither "Session.start" result) + pure sid + +parkFixture :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + Maybe Text -> + Eff es () +parkFixture sid key = do + now <- liftIO getCurrentTime + result <- + Session.awaitInput + AwaitInputData + { sessionId = sid, + reason = "approval", + correlationKey = key, + deadline = Nothing, + awaitedAt = now + } + void (liftEither "Session.awaitInput" result) + +resumeFixture :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + Maybe Text -> + Text -> + Eff es () +resumeFixture sid key input = do + now <- liftIO getCurrentTime + result <- + Session.resume + ResumeSessionData + { sessionId = sid, + correlationKey = key, + force = False, + input, + resumedAt = now + } + void (liftEither "Session.resume" result) + +getExisting :: + (IOE :> es, Store :> es) => + SessionId -> + Eff es SessionRow +getExisting sid = do + result <- Session.getById sid >>= liftEither "Session.getById" + case result of + Nothing -> liftIO (assertFailure ("missing session row " <> show (idText sid))) + Just row -> pure row + +readSessionEvents :: + (IOE :> es, Store :> es) => + SessionId -> + Eff es [SessionEvent] +readSessionEvents sid = do + recorded <- Vector.toList <$> readStreamForward (Stream.streamName (sessionStream sid)) (StreamVersion 0) 100 + traverse decodeRecorded recorded + where + decodeRecorded recorded = + case parseSessionEvent recorded.payload of + Left err -> liftIO (assertFailure ("parseSessionEvent: " <> show err)) + Right event -> pure event + +liftEither :: (Show e, IOE :> es) => String -> Either e a -> Eff es a +liftEither label = \case + Left err -> liftIO (assertFailure (label <> ": " <> show err)) + Right value -> pure value + +eventName :: SessionEvent -> Text +eventName = \case + SessionStarted {} -> "SessionStarted" + SessionAwaiting {} -> "SessionAwaiting" + SessionResumed {} -> "SessionResumed" + SessionCompleted {} -> "SessionCompleted" + SessionFailed {} -> "SessionFailed" + InteractiveSessionRecorded {} -> "InteractiveSessionRecorded" + TurnRecorded {} -> "TurnRecorded" + +testScope :: MemoryScope +testScope = ScopeGlobal (Namespace "kioku-test") diff --git a/test/Kioku/SessionLineageSpec.hs b/test/Kioku/SessionLineageSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/SessionLineageSpec.hs @@ -0,0 +1,195 @@ +module Kioku.SessionLineageSpec (tests) where + +import Control.Monad (void) +import Data.List (sort) +import Data.Text (Text) +import Data.Text.Encoding (encodeUtf8) +import Data.Time (UTCTime, addUTCTime, getCurrentTime) +import Effectful (Eff, IOE, liftIO, (:>)) +import Effectful.Error.Static (Error) +import Hasql.Transaction qualified as Tx +import Kioku.Api.Scope (MemoryScope (..), Namespace (..)) +import Kioku.App (runAppIO, withNoopAppEnv) +import Kioku.Id (SessionId, genSessionId, idText) +import Kioku.Migrations.TestSupport (withKiokuMigratedDatabase) +import Kioku.Session qualified as Session +import Kioku.Session.Domain (StartSessionData (..)) +import Kioku.Session.ReadModel (SessionRow (..)) +import Kiroku.Store.Connection (defaultConnectionSettings) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Transaction (runTransaction) +import Test.Tasty (TestTree, localOption, mkTimeout, testGroup) +import Test.Tasty.HUnit (assertEqual, assertFailure, testCase) + +tests :: TestTree +tests = + testGroup + "Session.Lineage" + [ testCase "child sessions record parent id and delegation depth" testDelegationLineage, + testGroup + "start rejects malformed lineage" + [ testCase "a session that is its own predecessor" (assertInvalidLineage selfPrevious), + testCase "a session that is its own parent" (assertInvalidLineage selfParent), + testCase "a negative delegation depth" (assertInvalidLineage negativeDepth), + testCase "a delegation depth past the cap" (assertInvalidLineage depthPastCap), + testCase "a delegated session at depth 0" (assertInvalidLineage parentAtDepthZero), + testCase "a root session at depth 1" (assertInvalidLineage rootAtDepthOne) + ], + -- A cycle cannot be built through the API any more, so this one manufactures corrupt + -- data directly. Before the CTE carried a path array, getChain looped until the + -- timeout fired; it now returns in milliseconds. + localOption (mkTimeout 10_000_000) $ + testCase "getChain terminates on a cyclic chain" testChainTerminatesOnCycle + ] + +data LineageResult = LineageResult + { children :: ![SessionRow], + childChain :: ![SessionRow] + } + +testDelegationLineage :: IO () +testDelegationLineage = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) \env -> do + parent <- genSessionId + child1 <- genSessionId + child2 <- genSessionId + now <- getCurrentTime + let child2Started = addUTCTime 2 now + result <- + runAppIO env do + startFixture parent "parent" Nothing Nothing 0 now + startFixture child1 "child-1" Nothing (Just parent) 1 (addUTCTime 1 now) + startFixture child2 "child-2" (Just child1) (Just parent) 1 child2Started + children <- Session.getDelegationChildren parent >>= liftEither "getDelegationChildren" + childChain <- Session.getChain child2 >>= liftEither "getChain" + pure LineageResult {children, childChain} + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right LineageResult {children, childChain} -> do + assertEqual "parent has exactly the two direct delegation children" [idText child1, idText child2] (rowSessionId <$> children) + assertEqual "children carry parentSessionId" [Just (idText parent), Just (idText parent)] (rowParentSessionId <$> children) + assertEqual "children carry depth 1" [1, 1] (rowDelegationDepth <$> children) + assertEqual "previous-session chain follows previousSessionId, not parentSessionId" [idText child1, idText child2] (rowSessionId <$> childChain) + +-- | A malformed-lineage case, built from the session's own id (so the self-reference cases +-- can point a field at it) and an unrelated id (so the depth cases can name a parent +-- without tripping the self-parent rule first). +type LineageCase = SessionId -> SessionId -> UTCTime -> StartSessionData + +baseStart :: LineageCase +baseStart sid _other startedAt = + StartSessionData + { sessionId = sid, + agentId = "test-agent", + focus = "delegation lineage", + scope = ScopeGlobal (Namespace "kioku-test"), + subjectRef = Nothing, + previousSessionId = Nothing, + parentSessionId = Nothing, + delegationDepth = 0, + startedAt + } + +selfPrevious, selfParent, negativeDepth, depthPastCap, parentAtDepthZero, rootAtDepthOne :: LineageCase +selfPrevious sid o t = (baseStart sid o t) {previousSessionId = Just sid} +selfParent sid o t = (baseStart sid o t) {parentSessionId = Just sid, delegationDepth = 1} +negativeDepth sid o t = (baseStart sid o t) {delegationDepth = -1} +depthPastCap sid o t = (baseStart sid o t) {parentSessionId = Just o, delegationDepth = 65} +parentAtDepthZero sid o t = (baseStart sid o t) {parentSessionId = Just o, delegationDepth = 0} +rootAtDepthOne sid o t = (baseStart sid o t) {delegationDepth = 1} + +assertInvalidLineage :: LineageCase -> IO () +assertInvalidLineage mkCommand = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) \env -> do + sid <- genSessionId + other <- genSessionId + now <- getCurrentTime + result <- runAppIO env (Session.start (mkCommand sid other now)) + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right (Left (Session.SessionInvalidLineage _)) -> pure () + Right other' -> + assertFailure ("expected SessionInvalidLineage, got " <> show other') + +startFixture :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + Text -> + Maybe SessionId -> + Maybe SessionId -> + Int -> + UTCTime -> + Eff es () +startFixture sid agent previous parent depth startedAt = do + result <- + Session.start + StartSessionData + { sessionId = sid, + agentId = agent, + focus = "delegation lineage", + scope = ScopeGlobal (Namespace "kioku-test"), + subjectRef = Nothing, + previousSessionId = previous, + parentSessionId = parent, + delegationDepth = depth, + startedAt + } + void (liftEither "Session.start" result) + +-- | Insert two sessions that name each other as predecessor, bypassing 'Session.start' +-- (which now refuses to create a cycle), then walk the chain. The assertion is simply that +-- the query returns at all — and returns each session exactly once, rather than looping. +testChainTerminatesOnCycle :: IO () +testChainTerminatesOnCycle = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) \env -> do + a <- genSessionId + b <- genSessionId + result <- + runAppIO env do + insertCyclicPair a b + Session.getChain a >>= liftEither "getChain" + case result of + Left storeErr -> assertFailure ("store error: " <> show storeErr) + Right chain -> do + assertEqual "the cyclic chain returns both sessions" 2 (length chain) + assertEqual + "each session appears exactly once" + (sort [idText a, idText b]) + (sort (rowSessionId <$> chain)) + +insertCyclicPair :: + (IOE :> es, Store :> es) => + SessionId -> + SessionId -> + Eff es () +insertCyclicPair a b = + runTransaction . Tx.sql . encodeUtf8 $ + "INSERT INTO kioku_sessions (session_id, agent_id, focus, namespace, delegation_depth, status, started_at, previous_session_id) VALUES " + <> "('" + <> idText a + <> "','t','f','kioku-test',0,'completed',NOW(),'" + <> idText b + <> "'),('" + <> idText b + <> "','t','f','kioku-test',0,'completed',NOW() - interval '1 second','" + <> idText a + <> "')" + +liftEither :: (Show e, IOE :> es) => String -> Either e a -> Eff es a +liftEither label = \case + Left err -> liftIO (assertFailure (label <> ": " <> show err)) + Right value -> pure value + +rowSessionId :: SessionRow -> Text +rowSessionId row = row.sessionId + +rowParentSessionId :: SessionRow -> Maybe Text +rowParentSessionId row = row.parentSessionId + +rowDelegationDepth :: SessionRow -> Int +rowDelegationDepth row = row.delegationDepth diff --git a/test/Kioku/TimerWorkerSpec.hs b/test/Kioku/TimerWorkerSpec.hs new file mode 100644 --- /dev/null +++ b/test/Kioku/TimerWorkerSpec.hs @@ -0,0 +1,296 @@ +{-# LANGUAGE DataKinds #-} + +module Kioku.TimerWorkerSpec + ( tests, + ) +where + +import Data.Aeson qualified as Aeson +import Data.Functor.Contravariant ((>$<)) +import Data.Int (Int64) +import Data.Text qualified as Text +import Data.Time (NominalDiffTime, addUTCTime, diffUTCTime) +import Data.UUID qualified as UUID +import Data.UUID.V4 qualified as UUIDv4 +import Effectful (Eff, IOE, (:>)) +import Effectful.Error.Static (Error) +import Hasql.Decoders qualified as D +import Hasql.Encoders qualified as E +import Hasql.Statement (Statement, preparable) +import Hasql.Transaction qualified as Tx +import Keiro.Timer (TimerId (..), TimerRequest (..), scheduleTimerTx) +import Kioku.Api.Scope (MemoryScope (..), Namespace (..), ScopeKind (..)) +import Kioku.App (AppEffects, AppEnv, runAppIO, withNoopAppEnv) +import Kioku.Distill.L1 (scopedScanCandidates) +import Kioku.Distill.L2 (l2SceneProcessManagerName) +import Kioku.Distill.Runtime (DistillRuntime (..), newDistillRuntime) +import Kioku.Distill.Timer (l1ExtractProcessManagerName) +import Kioku.Distill.Timer.Worker (drainKiokuTimers, runKiokuTimerWorkerOnce) +import Kioku.Id (SessionId, genSessionId, idText) +import Kioku.Migrations.TestSupport (withKiokuMigratedDatabase) +import Kioku.Prelude +import Kioku.Session qualified as Session +import Kioku.Session.Domain (StartSessionData (..)) +import Kiroku.Store.Connection (defaultConnectionSettings) +import Kiroku.Store.Effect (Store) +import Kiroku.Store.Effect.Resource (KirokuStoreResource) +import Kiroku.Store.Error (StoreError) +import Kiroku.Store.Transaction (runTransaction) +import Shikumi.Error (ShikumiError (..)) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Timer worker" + [ testCase "permanent failure dead-letters the timer" testPermanentFailureDeadLetters, + testCase "transient failure reschedules with backoff" testTransientFailureReschedules, + testCase "unknown process manager requeues with a long delay" testUnknownProcessManagerRequeues, + testCase "the attempt ceiling dead-letters" testAttemptCeilingDeadLetters, + testCase "success marks the timer fired" testSuccessMarksFired, + testCase "drain processes every due timer in one pass" testDrainProcessesAllDueTimers + ] + +-- | A correlation id that is not a session id can never become one. It used to +-- be marked fired — a fake success that lost the distillation silently. +testPermanentFailureDeadLetters :: Assertion +testPermanentFailureDeadLetters = + withTimerEnv \env rt -> do + timerId <- freshTimerId + row <- runOrFail env do + scheduleTestTimer timerId l1ExtractProcessManagerName "not-a-session-id" Aeson.Null (-1) + fireOnce rt + fetchTimer timerId + row.status @?= "dead" + assertBool + ("last_error names the correlation id, got: " <> show row.lastError) + (maybe False (Text.isInfixOf "correlation id") row.lastError) + +-- | A failing LLM extraction is worth retrying — but on a schedule, and only +-- until the ceiling. The timer must land back in @scheduled@ with @fire_at@ +-- pushed out, not sit in @firing@ waiting on keiro's 300-second stale requeue. +testTransientFailureReschedules :: Assertion +testTransientFailureReschedules = + withTimerEnv \env rt -> do + timerId <- freshTimerId + sid <- genSessionId + let failing = rt {runExtract = \_ -> pure (Left (ProviderFailure "the model is down"))} + before <- getCurrentTime + row <- runOrFail env do + startFixtureSession sid + scheduleTestTimer timerId l1ExtractProcessManagerName (idText sid) Aeson.Null (-1) + fireOnce failing + fetchTimer timerId + row.status @?= "scheduled" + -- keiro increments attempts at claim time, so one claim means attempts = 1, + -- and the first backoff step is 30s. + row.attempts @?= 1 + assertDelayNear "first retry" before 30 row.fireAt + +-- | keiro's claimDueTimer claims the earliest due timer regardless of process +-- manager, so a timer no handler owns cannot be left alone — it must be put +-- back, or it starves every other timer behind it forever. +testUnknownProcessManagerRequeues :: Assertion +testUnknownProcessManagerRequeues = + withTimerEnv \env rt -> do + timerId <- freshTimerId + before <- getCurrentTime + row <- runOrFail env do + scheduleTestTimer timerId "kioku-nonexistent" "whatever" Aeson.Null (-1) + fireOnce rt + fetchTimer timerId + row.status @?= "scheduled" + row.attempts @?= 1 + assertDelayNear "unknown-PM requeue" before 600 row.fireAt + +-- | The requeue above is bounded: an orphaned timer eventually dies visibly +-- instead of cycling forever. keiro applies the ceiling at claim time. +testAttemptCeilingDeadLetters :: Assertion +testAttemptCeilingDeadLetters = + withTimerEnv \env rt -> do + timerId <- freshTimerId + row <- runOrFail env do + scheduleTestTimer timerId "kioku-nonexistent" "whatever" Aeson.Null (-1) + forceAttempts timerId 8 + fireOnce rt + fetchTimer timerId + row.status @?= "dead" + assertBool + ("last_error mentions the ceiling, got: " <> show row.lastError) + (maybe False (Text.isInfixOf "attempt ceiling") row.lastError) + +-- | The happy path: an L2 timer for a scope with no memories regenerates +-- nothing, succeeds without calling the LLM, and is marked fired with the timer's +-- own id as the marker event id. +testSuccessMarksFired :: Assertion +testSuccessMarksFired = + withTimerEnv \env rt -> do + timerId@(TimerId timerUuid) <- freshTimerId + row <- runOrFail env do + scheduleTestTimer + timerId + l2SceneProcessManagerName + "rei/intention/empty" + (Aeson.object ["scope" Aeson..= emptyScope]) + (-1) + fireOnce rt + fetchTimer timerId + row.status @?= "fired" + row.firedEventId @?= Just (UUID.toText timerUuid) + +-- | The old loop slept between every timer, so N due timers took N poll +-- intervals. Draining processes them all in one pass. +testDrainProcessesAllDueTimers :: Assertion +testDrainProcessesAllDueTimers = + withTimerEnv \env rt -> do + timerIds <- traverse (const freshTimerId) [1 :: Int, 2, 3] + (processed, rows) <- runOrFail env do + forM_ timerIds \timerId -> + scheduleTestTimer timerId l1ExtractProcessManagerName "not-a-session-id" Aeson.Null (-1) + processed <- drainKiokuTimers Nothing rt (scopedScanCandidates 5) + rows <- traverse fetchTimer timerIds + pure (processed, rows) + processed @?= 3 + fmap (.status) rows @?= ["dead", "dead", "dead"] + +fireOnce :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + DistillRuntime -> + Eff es () +fireOnce rt = do + now <- liftIO getCurrentTime + void (runKiokuTimerWorkerOnce Nothing rt (scopedScanCandidates 5) now) + +-- | Schedule a timer @offset@ seconds from now (negative means already due). +scheduleTestTimer :: + (IOE :> es, Store :> es) => + TimerId -> + Text -> + Text -> + Aeson.Value -> + NominalDiffTime -> + Eff es () +scheduleTestTimer timerId processManagerName correlationId payload offset = do + now <- liftIO getCurrentTime + runTransaction $ + scheduleTimerTx + TimerRequest + { timerId, + processManagerName, + correlationId, + fireAt = addUTCTime offset now, + payload + } + +startFixtureSession :: + (IOE :> es, KirokuStoreResource :> es, Store :> es, Error StoreError :> es) => + SessionId -> + Eff es () +startFixtureSession sid = do + now <- liftIO getCurrentTime + started <- + Session.start + StartSessionData + { sessionId = sid, + agentId = "test-agent", + focus = "timer worker spec", + scope = emptyScope, + subjectRef = Nothing, + previousSessionId = Nothing, + parentSessionId = Nothing, + delegationDepth = 0, + startedAt = now + } + void (liftIO (expectRight "Session.start" started)) + +-- | Drive the row to the brink of the ceiling so the next claim trips it. +forceAttempts :: (Store :> es) => TimerId -> Int -> Eff es () +forceAttempts (TimerId uuid) attempts = + runTransaction (Tx.statement (uuid, fromIntegral @Int @Int64 attempts) forceAttemptsStmt) + +forceAttemptsStmt :: Statement (UUID.UUID, Int64) () +forceAttemptsStmt = + preparable + """ + UPDATE keiro.keiro_timers + SET attempts = $2, + fire_at = now() - interval '1 second' + WHERE timer_id = $1 + """ + ( ((\(u, _) -> u) >$< E.param (E.nonNullable E.uuid)) + <> ((\(_, a) -> a) >$< E.param (E.nonNullable E.int8)) + ) + D.noResult + +data TimerStateRow = TimerStateRow + { status :: !Text, + attempts :: !Int, + fireAt :: !UTCTime, + lastError :: !(Maybe Text), + firedEventId :: !(Maybe Text) + } + deriving stock (Generic, Eq, Show) + +fetchTimer :: (Store :> es) => TimerId -> Eff es TimerStateRow +fetchTimer (TimerId uuid) = + runTransaction (Tx.statement uuid selectTimerStateStmt) + +selectTimerStateStmt :: Statement UUID.UUID TimerStateRow +selectTimerStateStmt = + preparable + """ + SELECT status, attempts, fire_at, last_error, fired_event_id::text + FROM keiro.keiro_timers + WHERE timer_id = $1 + """ + (E.param (E.nonNullable E.uuid)) + (D.singleRow timerStateDecoder) + +timerStateDecoder :: D.Row TimerStateRow +timerStateDecoder = + TimerStateRow + <$> D.column (D.nonNullable D.text) + <*> (fromIntegral @Int64 @Int <$> D.column (D.nonNullable D.int8)) + <*> D.column (D.nonNullable D.timestamptz) + <*> D.column (D.nullable D.text) + <*> D.column (D.nullable D.text) + +-- | The delay is measured against a clock read taken before the pass, so the +-- window has to absorb the pass's own duration. A generous window still +-- distinguishes 30s from 600s, which is the distinction under test. +assertDelayNear :: String -> UTCTime -> NominalDiffTime -> UTCTime -> Assertion +assertDelayNear label before expected actual = + assertBool + ( label + <> ": expected fire_at about " + <> show expected + <> "s out, but it was " + <> show delay + <> "s out" + ) + (delay >= expected - 5 && delay <= expected + 15) + where + delay = actual `diffUTCTime` before + +emptyScope :: MemoryScope +emptyScope = + ScopeEntity (Namespace "rei") (ScopeKind "intention") "intention_timer_worker_spec" + +freshTimerId :: IO TimerId +freshTimerId = TimerId <$> UUIDv4.nextRandom + +withTimerEnv :: (AppEnv -> DistillRuntime -> IO ()) -> Assertion +withTimerEnv action = + withKiokuMigratedDatabase \connStr -> + withNoopAppEnv (defaultConnectionSettings connStr) \env -> do + rt <- newDistillRuntime + action env rt + +runOrFail :: AppEnv -> Eff AppEffects a -> IO a +runOrFail env action = runAppIO env action >>= expectRight "runAppIO" + +expectRight :: (Show e) => String -> Either e a -> IO a +expectRight label = \case + Left err -> assertFailure (label <> " failed: " <> show err) + Right value -> pure value diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,36 @@ +module Main where + +import Kioku.AwaitingSpec qualified as AwaitingSpec +import Kioku.DistillSpec qualified as DistillSpec +import Kioku.EmbeddingWorkerSpec qualified as EmbeddingWorkerSpec +import Kioku.IdempotencySpec qualified as IdempotencySpec +import Kioku.ReadModelReconcileSpec qualified as ReadModelReconcileSpec +import Kioku.RecallSpec qualified as RecallSpec +import Kioku.RecallSqlSpec qualified as RecallSqlSpec +import Kioku.ReiCompatSpec qualified as ReiCompatSpec +import Kioku.SchemaSpec qualified as SchemaSpec +import Kioku.ScopeIdentitySpec qualified as ScopeIdentitySpec +import Kioku.SessionInvariantsSpec qualified as SessionInvariantsSpec +import Kioku.SessionLineageSpec qualified as SessionLineageSpec +import Kioku.TimerWorkerSpec qualified as TimerWorkerSpec +import Test.Tasty (defaultMain, testGroup) + +main :: IO () +main = + defaultMain $ + testGroup + "kioku" + [ AwaitingSpec.tests, + ReiCompatSpec.tests, + IdempotencySpec.tests, + ReadModelReconcileSpec.tests, + RecallSpec.tests, + RecallSqlSpec.tests, + SchemaSpec.tests, + ScopeIdentitySpec.tests, + SessionInvariantsSpec.tests, + SessionLineageSpec.tests, + EmbeddingWorkerSpec.tests, + TimerWorkerSpec.tests, + DistillSpec.tests + ]