shikumi-trace (empty) → 0.1.0.0
raw patch · 14 files changed
+2073/−0 lines, 14 filesdep +QuickCheckdep +aesondep +baikai
Dependencies added: QuickCheck, aeson, baikai, base, bytestring, containers, directory, effectful, generic-lens, lens, scientific, shikumi, shikumi-cache, shikumi-trace, tasty, tasty-hunit, tasty-quickcheck, temporary, text, time, vector
Files
- CHANGELOG.md +10/−0
- app/Main.hs +14/−0
- shikumi-trace.cabal +104/−0
- src/Shikumi/Trace.hs +405/−0
- src/Shikumi/Trace/Demo.hs +169/−0
- src/Shikumi/Trace/Feedback.hs +78/−0
- src/Shikumi/Trace/Internal/Spike.hs +142/−0
- src/Shikumi/Trace/Node.hs +126/−0
- src/Shikumi/Trace/Program.hs +186/−0
- src/Shikumi/Trace/Replay.hs +95/−0
- src/Shikumi/Trace/ResponseJSON.hs +17/−0
- src/Shikumi/Trace/Store.hs +91/−0
- test/Main.hs +533/−0
- test/TraceFixtures.hs +103/−0
+ CHANGELOG.md view
@@ -0,0 +1,10 @@+# Changelog++## Unreleased++## 0.1.0.0 - 2026-06-13++### Added++- Initial Hackage release of hierarchical tracing and deterministic replay for shikumi.+- Trace spans, trace trees, program-node tracing, replay, trace storage, feedback logs, rendering, and a trace demo executable.
+ app/Main.hs view
@@ -0,0 +1,14 @@+-- | The @shikumi-trace-demo@ executable (EP-7, M5). Delegates to the testable+-- core 'Shikumi.Trace.Demo.demoMain':+--+-- @+-- cabal run shikumi-trace-demo -- live: prints the tree, writes trace.json+-- SHIKUMI_OFFLINE=1 cabal run shikumi-trace-demo -- --replay trace.json+-- @+module Main (main) where++import Shikumi.Trace.Demo (demoMain)+import System.Environment (getArgs)++main :: IO ()+main = getArgs >>= demoMain
+ shikumi-trace.cabal view
@@ -0,0 +1,104 @@+cabal-version: 3.4+name: shikumi-trace+version: 0.1.0.0+synopsis:+ Hierarchical tracing, observability, and deterministic replay for shikumi (EP-7)++category: AI+description:+ Hierarchical (nested) tracing and deterministic offline replay for shikumi.+ Running a program inside the @Trace@ effect produces a /tree/ of spans (program+ / module / combinator / LM-call), captured by interposing on EP-1's @LLM@+ effect — so each provider call becomes a leaf span carrying model, latency,+ tokens, cost, and the EP-6 content-addressed cache key. The tree pretty-prints+ to the terminal and serializes to a stable JSON trace file; @runLLMReplay@ then+ re-runs the same program with the network disabled, serving every call from the+ recording (and raising a precise @ReplayDivergence@ for any unrecorded call).++license: BSD-3-Clause+author: Nadeem Bitar+maintainer: nadeem@gmail.com+build-type: Simple+extra-doc-files: CHANGELOG.md++common common-options+ ghc-options:+ -Wall -Wcompat -Widentities -Wincomplete-uni-patterns+ -Wincomplete-record-updates -Wredundant-constraints+ -fhide-source-paths -Wmissing-export-lists -Wpartial-fields+ -Wmissing-deriving-strategies++ default-language: GHC2024+ default-extensions:+ DeriveAnyClass+ DuplicateRecordFields+ OverloadedLabels+ OverloadedStrings++library+ import: common-options+ hs-source-dirs: src+ exposed-modules:+ Shikumi.Trace+ Shikumi.Trace.Demo+ Shikumi.Trace.Feedback+ Shikumi.Trace.Internal.Spike+ Shikumi.Trace.Node+ Shikumi.Trace.Program+ Shikumi.Trace.Replay+ Shikumi.Trace.ResponseJSON+ Shikumi.Trace.Store++ build-depends:+ , aeson+ , baikai+ , base >=4.20 && <5+ , bytestring+ , containers+ , directory+ , effectful+ , generic-lens+ , lens ^>=5.3+ , scientific+ , shikumi ^>=0.1.0.0+ , shikumi-cache ^>=0.1.0.0+ , text ^>=2.1+ , time+ , vector++executable shikumi-trace-demo+ import: common-options+ hs-source-dirs: app+ main-is: Main.hs+ ghc-options: -threaded -with-rtsopts=-N+ build-depends:+ , base+ , shikumi-trace ^>=0.1.0.0++test-suite shikumi-trace-test+ import: common-options+ type: exitcode-stdio-1.0+ hs-source-dirs: test+ main-is: Main.hs+ ghc-options: -threaded -with-rtsopts=-N+ other-modules: TraceFixtures+ build-depends:+ , aeson+ , baikai+ , base+ , bytestring+ , containers+ , effectful+ , generic-lens+ , lens ^>=5.3+ , QuickCheck+ , shikumi ^>=0.1.0.0+ , shikumi-cache ^>=0.1.0.0+ , shikumi-trace ^>=0.1.0.0+ , tasty+ , tasty-hunit+ , tasty-quickcheck+ , temporary+ , text ^>=2.1+ , time+ , vector
+ src/Shikumi/Trace.hs view
@@ -0,0 +1,405 @@+{-# LANGUAGE GADTs #-}+{-# LANGUAGE TypeFamilies #-}++-- | Hierarchical tracing for shikumi (EP-7, M1).+--+-- The 'Trace' effect produces a /tree/ of timed, attributed spans for a program+-- run. A span is a node of one of four 'SpanKind's — a whole program, a module+-- such as @predict@, a combinator such as @Pipeline@, or a single LM call. Each+-- span records its parent, so nesting @withSpan@s builds a real tree even though+-- neither baikai nor EP-1's @LLM@ effect has any parent\/child concept.+--+-- The hierarchy is shikumi's: 'runTrace' keeps a stack of span ids, and 'withSpan'+-- pushes\/pops. LM-call leaves are captured automatically by 'tracedLLM', which+-- __interposes__ on EP-1's @LLM@ effect (the same seam EP-6's @cachedLLM@ uses):+-- on each 'Shikumi.LLM.complete' it opens an 'LlmCallSpan' under the active span+-- and fills its attributes — model, provider, latency, tokens, cost, tool calls,+-- the recorded response JSON, and the EP-6 content-addressed 'Shikumi.Cache.Key.cacheKey'+-- (integration point #7) — from the returned 'Baikai.Response'.+--+-- 'renderTree' pretty-prints the tree; 'Shikumi.Trace.Store' serializes it and+-- 'Shikumi.Trace.Replay' replays it offline.+module Shikumi.Trace+ ( -- * Span and tree types+ SpanKind (..),+ SpanId (..),+ ToolCallRecord (..),+ SpanAttrs (..),+ emptyAttrs,+ Span (..),+ TraceTree (..),+ childrenOf,++ -- * Node identity (re-exported from "Shikumi.Trace.Node")+ NodePath (..),++ -- * The effect+ Trace,+ withSpan,+ currentSpanId,+ bumpRetry,+ recordToolCall,+ annotateSpan,++ -- * Interpreters and capture+ runTrace,+ tracedLLM,+ llmLabel,+ llmAttrs,++ -- * Rendering+ renderTree,+ )+where++import Baikai+ ( AssistantContent (..),+ Context,+ Model,+ Options,+ Response,+ flattenAssistantBlocks,+ )+import Control.Lens ((^.))+import Data.Aeson (FromJSON, FromJSONKey, ToJSON, ToJSONKey, Value, toJSON)+import Data.Generics.Labels ()+import Data.List (sortOn)+import Data.Map.Strict (Map)+import Data.Map.Strict qualified as Map+import Data.Maybe (fromMaybe)+import Data.Scientific (Scientific)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Time.Clock (UTCTime, diffUTCTime)+import Data.Vector qualified as V+import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>))+import Effectful.Dispatch.Dynamic (interpose, interpret, localSeqUnlift, send)+import Effectful.Exception (bracket)+import Effectful.Prim (Prim)+import Effectful.Prim.IORef+ ( IORef,+ atomicModifyIORef',+ modifyIORef',+ newIORef,+ readIORef,+ )+import GHC.Generics (Generic)+import Numeric.Natural (Natural)+import Shikumi.Cache.Key (CacheKey (..), requestToCanonicalValue)+import Shikumi.Cache.Key qualified as Key+import Shikumi.Effect.Time (Time, getCurrentTime)+import Shikumi.LLM (LLM (..), complete, stream)+import Shikumi.Trace.Node (NodePath (..))+import Shikumi.Trace.ResponseJSON ()++-- ---------------------------------------------------------------------------+-- Span and tree types+-- ---------------------------------------------------------------------------++-- | What sort of work a span records.+data SpanKind = ProgramSpan | ModuleSpan | CombinatorSpan | LlmCallSpan+ deriving stock (Eq, Show, Generic)+ deriving anyclass (ToJSON, FromJSON)++-- | A span identifier (a fresh @"span-N"@ string allocated by 'runTrace').+newtype SpanId = SpanId Text+ deriving stock (Eq, Ord, Show, Generic)+ deriving newtype (ToJSON, FromJSON, ToJSONKey, FromJSONKey)++-- | A tool call observed on an LM response: the tool name and its JSON arguments.+data ToolCallRecord = ToolCallRecord+ { name :: !Text,+ arguments :: !Value+ }+ deriving stock (Eq, Show, Generic)+ deriving anyclass (ToJSON, FromJSON)++-- | The spec-mandated attribute bag a span can carry. Most fields are populated+-- only on 'LlmCallSpan' nodes; the structural nodes (program\/module\/combinator)+-- leave them empty.+data SpanAttrs = SpanAttrs+ { -- | model id, for LM-call spans+ model :: !(Maybe Text),+ provider :: !(Maybe Text),+ -- | the request as canonical JSON (system, messages, tools, options)+ prompt :: !(Maybe Value),+ -- | the response payload as JSON (for replay + inspection)+ response :: !(Maybe Value),+ latencyMs :: !(Maybe Integer),+ inputTokens :: !(Maybe Natural),+ outputTokens :: !(Maybe Natural),+ costUsd :: !(Maybe Scientific),+ -- | how many times this span retried its body+ retries :: !Int,+ -- | tool calls observed on the response+ toolCalls :: ![ToolCallRecord],+ -- | the EP-6 content-addressed key, present on LM-call spans+ cacheKey :: !(Maybe Text),+ -- | the structural path of the @Program@ node that issued this span's LM call+ -- (EP-16). Present only on model-call spans produced by @runProgramTraced@;+ -- 'Nothing' for spans opened by bare 'withSpan' or a non-node-correlated run.+ nodePath :: !(Maybe NodePath)+ }+ deriving stock (Eq, Show, Generic)+ deriving anyclass (ToJSON, FromJSON)++-- | A span with everything empty: the base every node starts from.+emptyAttrs :: SpanAttrs+emptyAttrs =+ SpanAttrs+ { model = Nothing,+ provider = Nothing,+ prompt = Nothing,+ response = Nothing,+ latencyMs = Nothing,+ inputTokens = Nothing,+ outputTokens = Nothing,+ costUsd = Nothing,+ retries = 0,+ toolCalls = [],+ cacheKey = Nothing,+ nodePath = Nothing+ }++-- | One node of the trace tree.+data Span = Span+ { spanId :: !SpanId,+ parent :: !(Maybe SpanId),+ kind :: !SpanKind,+ label :: !Text,+ startedAt :: !UTCTime,+ endedAt :: !(Maybe UTCTime),+ attrs :: !SpanAttrs+ }+ deriving stock (Eq, Show, Generic)+ deriving anyclass (ToJSON, FromJSON)++-- | A finished trace tree: a flat 'Map' of spans plus the root id. The nesting is+-- reconstructed from each span's 'parent' pointer ('childrenOf').+data TraceTree = TraceTree+ { root :: !SpanId,+ spans :: !(Map SpanId Span)+ }+ deriving stock (Eq, Show, Generic)+ deriving anyclass (ToJSON, FromJSON)++-- | The children of a span, in creation order (sorted by start time, ties broken+-- by span id).+childrenOf :: TraceTree -> SpanId -> [SpanId]+childrenOf t sid =+ map spanId $+ sortOn (\s -> (startedAt s, spanId s)) $+ [s | s <- Map.elems (spans t), parent s == Just sid]++-- ---------------------------------------------------------------------------+-- The effect+-- ---------------------------------------------------------------------------++-- | The hierarchical tracing effect.+data Trace :: Effect where+ WithSpan :: SpanKind -> Text -> m a -> Trace m a+ CurrentSpanId :: Trace m (Maybe SpanId)+ BumpRetry :: Trace m ()+ RecordToolCall :: ToolCallRecord -> Trace m ()+ AnnotateSpan :: (SpanAttrs -> SpanAttrs) -> Trace m ()++type instance DispatchOf Trace = 'Dynamic++-- | Open a span of the given kind and label around an inner computation: it is+-- pushed as a child of the currently-active span, timed, and recorded on exit.+withSpan :: (Trace :> es) => SpanKind -> Text -> Eff es a -> Eff es a+withSpan k lbl act = send (WithSpan k lbl act)++-- | The id of the currently-active span, if any.+currentSpanId :: (Trace :> es) => Eff es (Maybe SpanId)+currentSpanId = send CurrentSpanId++-- | Increment the active span's retry counter.+bumpRetry :: (Trace :> es) => Eff es ()+bumpRetry = send BumpRetry++-- | Record a tool call on the active span.+recordToolCall :: (Trace :> es) => ToolCallRecord -> Eff es ()+recordToolCall = send . RecordToolCall++-- | Apply a function to the active span's attributes (used by 'tracedLLM' to fill+-- an LM-call span after the response arrives).+annotateSpan :: (Trace :> es) => (SpanAttrs -> SpanAttrs) -> Eff es ()+annotateSpan = send . AnnotateSpan++-- ---------------------------------------------------------------------------+-- The interpreter+-- ---------------------------------------------------------------------------++-- | Mutable building state for one trace, all in 'IORef's.+data TraceState = TraceState+ { tsCounter :: !(IORef Int),+ tsStack :: !(IORef [SpanId]),+ tsSpans :: !(IORef (Map SpanId Span)),+ tsRoot :: !(IORef (Maybe SpanId))+ }++-- | Run a traced computation, returning its result and the finished tree.+--+-- The building state lives in 'IORef's reached through the 'Prim' effect (so no+-- open-ended @IOE@ is needed here — only in-process mutation), and span+-- timestamps come from shikumi's own 'Time' effect. Both are discharged at the+-- program edge by 'runPrim' and 'runTime'.+runTrace :: (Prim :> es, Time :> es) => Eff (Trace : es) a -> Eff es (a, TraceTree)+runTrace act = do+ st <- newTraceState+ a <-+ interpret+ ( \env -> \case+ CurrentSpanId -> safeHead <$> readIORef (tsStack st)+ BumpRetry -> modifyActive st (\a' -> a' {retries = retries a' + 1})+ RecordToolCall tc -> modifyActive st (\a' -> a' {toolCalls = toolCalls a' ++ [tc]})+ AnnotateSpan f -> modifyActive st f+ WithSpan k lbl inner ->+ bracket+ (openSpan st k lbl)+ (closeSpan st)+ (\_ -> localSeqUnlift env (\unlift -> unlift inner))+ )+ act+ tree <- freezeTree st+ pure (a, tree)++newTraceState :: (Prim :> es) => Eff es TraceState+newTraceState =+ TraceState <$> newIORef 0 <*> newIORef [] <*> newIORef Map.empty <*> newIORef Nothing++-- | Allocate a fresh span as a child of the current top-of-stack, insert it+-- (open, 'endedAt' = 'Nothing'), push it, and record it as the root if it is the+-- first parentless span.+openSpan :: (Prim :> es, Time :> es) => TraceState -> SpanKind -> Text -> Eff es SpanId+openSpan st k lbl = do+ n <- atomicModifyIORef' (tsCounter st) (\i -> (i + 1, i + 1))+ let sid = SpanId ("span-" <> T.pack (show n))+ par <- safeHead <$> readIORef (tsStack st)+ now <- getCurrentTime+ let s = Span sid par k lbl now Nothing emptyAttrs+ modifyIORef' (tsSpans st) (Map.insert sid s)+ modifyIORef' (tsStack st) (sid :)+ case par of+ Nothing -> modifyIORef' (tsRoot st) (Just . fromMaybe sid)+ Just _ -> pure ()+ pure sid++-- | Close a span: stamp its 'endedAt' and pop it off the stack.+closeSpan :: (Prim :> es, Time :> es) => TraceState -> SpanId -> Eff es ()+closeSpan st sid = do+ now <- getCurrentTime+ modifyIORef' (tsSpans st) (Map.adjust (\s -> s {endedAt = Just now}) sid)+ modifyIORef' (tsStack st) (drop 1)++-- | Apply a function to the active span's attributes (a no-op with no active span).+modifyActive :: (Prim :> es) => TraceState -> (SpanAttrs -> SpanAttrs) -> Eff es ()+modifyActive st f = do+ stk <- readIORef (tsStack st)+ case stk of+ (sid : _) -> modifyIORef' (tsSpans st) (Map.adjust (\s -> s {attrs = f (attrs s)}) sid)+ [] -> pure ()++-- | Freeze the building state into an immutable 'TraceTree'.+freezeTree :: (Prim :> es) => TraceState -> Eff es TraceTree+freezeTree st = do+ sp <- readIORef (tsSpans st)+ r <- readIORef (tsRoot st)+ pure (TraceTree (fromMaybe (SpanId "") r) sp)++safeHead :: [a] -> Maybe a+safeHead [] = Nothing+safeHead (x : _) = Just x++-- ---------------------------------------------------------------------------+-- LM-call capture (interpose on the LLM effect)+-- ---------------------------------------------------------------------------++-- | Capture every LM call as a leaf 'LlmCallSpan' under the active span. Interpose+-- on EP-1's @LLM@ effect: open a span, delegate to the underlying handler, then+-- fill the span's attributes from the returned 'Response' (and the request). The+-- streaming op is wrapped in a span but its attributes are left empty (streams+-- carry the same data incrementally; the demo/replay path uses 'complete').+tracedLLM :: (Trace :> es, LLM :> es) => Eff es a -> Eff es a+tracedLLM = interpose $ \_ -> \case+ Complete m c o -> withSpan LlmCallSpan (llmLabel m) $ do+ resp <- complete m c o+ annotateSpan (const (llmAttrs m c o resp))+ pure resp+ Stream m c o -> withSpan LlmCallSpan (llmLabel m) (stream m c o)++-- | The label for an LM-call span: @provider/model-id@.+llmLabel :: Model -> Text+llmLabel m = (m ^. #provider) <> "/" <> (m ^. #modelId)++-- | Build the attribute bag for a completed LM call from the request and response.+llmAttrs :: Model -> Context -> Options -> Response -> SpanAttrs+llmAttrs m c o resp =+ emptyAttrs+ { model = Just (m ^. #modelId),+ provider = Just (m ^. #provider),+ prompt = Just (requestToCanonicalValue m c o),+ response = Just (toJSON resp),+ latencyMs = Just (resp ^. #latencyMs),+ inputTokens = Just (resp ^. #message . #usage . #inputTokens),+ outputTokens = Just (resp ^. #message . #usage . #outputTokens),+ costUsd = Just (realToFrac (resp ^. #message . #usage . #cost . #usd :: Rational)),+ toolCalls = toolCallsOf resp,+ cacheKey = Just (unCacheKey (Key.cacheKey m c o))+ }++-- | Extract the tool calls from a response's assistant content blocks.+toolCallsOf :: Response -> [ToolCallRecord]+toolCallsOf resp =+ [ ToolCallRecord (tc ^. #name) (tc ^. #arguments)+ | AssistantToolCall tc <- V.toList (flattenAssistantBlocks resp)+ ]++-- ---------------------------------------------------------------------------+-- Rendering+-- ---------------------------------------------------------------------------++-- | Pretty-print a trace tree as an indented outline, one line per span: kind,+-- label, wall-clock duration, and (for LM-call spans) tokens and cost.+renderTree :: TraceTree -> Text+renderTree t+ | Map.null (spans t) = "(empty trace)\n"+ | otherwise = T.concat (go (0 :: Int) (root t))+ where+ go depth sid = case Map.lookup sid (spans t) of+ Nothing -> []+ Just s -> line depth s : concatMap (go (depth + 1)) (childrenOf t sid)+ line depth s =+ T.replicate (depth * 2) " "+ <> kindTag (kind s)+ <> " "+ <> label s+ <> durationOf s+ <> llmStats s+ <> "\n"+ durationOf s = case endedAt s of+ Just e ->+ let ms = round (realToFrac (diffUTCTime e (startedAt s)) * 1000 :: Double) :: Integer+ in " " <> tshow ms <> "ms"+ Nothing -> ""+ llmStats s+ | kind s == LlmCallSpan =+ let a = attrs s+ in " in="+ <> maybe "?" tshow (inputTokens a)+ <> " out="+ <> maybe "?" tshow (outputTokens a)+ <> " $"+ <> maybe "0" tshow (costUsd a)+ | otherwise = ""++kindTag :: SpanKind -> Text+kindTag = \case+ ProgramSpan -> "program"+ ModuleSpan -> "module"+ CombinatorSpan -> "combinator"+ LlmCallSpan -> "llm-call"++tshow :: (Show a) => a -> Text+tshow = T.pack . show
+ src/Shikumi/Trace/Demo.hs view
@@ -0,0 +1,169 @@+-- | The end-to-end demo (EP-7, M5): a two-stage pipeline that drafts a one-line+-- summary of a fixed article and then critiques the draft, tying together the+-- whole story — live trace → pretty-print → persist → deterministic offline+-- replay with identical output and zero provider calls.+--+-- The pipeline ('demoPipeline') is the shared core used by both the+-- @shikumi-trace-demo@ executable ('demoMain') and the @-p e2e@ acceptance test.+-- It is provider-agnostic: it issues 'Shikumi.LLM.complete' calls through the+-- @LLM@ effect and nests its stages in 'withSpan's, so the live run captures a+-- tree and the replay run is byte-identical.+--+-- For reproducibility the demo uses a deterministic in-process stub+-- ('demoResponder') rather than a real provider — no API keys, no network — so a+-- novice can run it anywhere and replay is well-defined. (The plan sketched an+-- optional live-provider path; the canned stub is the hermetic, always-runnable+-- choice. Replaying never contacts a provider regardless: 'runLLMReplay' has no+-- registry.)+module Shikumi.Trace.Demo+ ( demoArticle,+ demoModel,+ demoPipeline,+ demoResponder,+ runStubLLM,+ demoMain,+ )+where++import Baikai+ ( Api (Custom),+ AssistantContent (..),+ Context,+ Message (UserMessage),+ Model,+ Options,+ Response,+ TextContent (..),+ UserContent (UserText),+ user,+ _Context,+ _Model,+ _Options,+ _Response,+ _TextContent,+ )+import Control.Lens ((&), (.~), (^.))+import Data.Generics.Labels ()+import Data.Text (Text)+import Data.Text qualified as T+import Data.Text.IO qualified as TIO+import Data.Vector qualified as V+import Effectful (Eff, runEff, (:>))+import Effectful.Dispatch.Dynamic (interpret)+import Effectful.Prim (runPrim)+import Shikumi.Effect.Time (runTime)+import Shikumi.LLM (LLM (..), complete)+import Shikumi.Trace+ ( SpanKind (ModuleSpan, ProgramSpan),+ Trace,+ renderTree,+ runTrace,+ tracedLLM,+ withSpan,+ )+import Shikumi.Trace.Replay (runLLMReplay)+import Shikumi.Trace.Store (readTraceFile, replayIndex, writeTraceFile)+import System.Environment (lookupEnv)++-- | The fixed input article the demo summarizes.+demoArticle :: Text+demoArticle = "Shikumi turns LM calls into typed, traceable, replayable programs."++-- | The demo model's routing identity (LM-call spans read @stub/stub-model@).+demoModel :: Model+demoModel =+ _Model+ & #provider .~ "stub"+ & #modelId .~ "stub-model"+ & #api .~ Custom "stub"++-- | The two-stage pipeline: draft a one-line summary, then critique the draft.+-- Stage two's request depends on stage one's output, so replay must reproduce+-- stage one exactly to reach stage two. Returns the critique.+demoPipeline :: (Trace :> es, LLM :> es) => Text -> Eff es Text+demoPipeline article =+ withSpan ProgramSpan "summarize-and-critique" $ do+ draft <-+ withSpan ModuleSpan "predict:Draft" $+ textOf <$> complete demoModel (ctxFor ("draft: " <> article)) opts+ withSpan ModuleSpan "predict:Critique" $+ textOf <$> complete demoModel (ctxFor ("critique: " <> draft)) opts++-- | The deterministic stub: a draft request yields a fixed summary; a critique+-- request returns a fixed verdict. Deterministic per request, so replay is+-- well-defined.+demoResponder :: Context -> Response+demoResponder c =+ if "draft:" `T.isPrefixOf` lastUserText c+ then mkResponse "Shikumi makes LM programs typed and replayable."+ else mkResponse "Accurate and concise; ship it."++-- | A base @LLM@ interpreter driving the demo from 'demoResponder'.+runStubLLM :: Eff (LLM : es) a -> Eff es a+runStubLLM = interpret $ \_ -> \case+ Complete _ c _ -> pure (demoResponder c)+ Stream {} -> pure []++-- ---------------------------------------------------------------------------+-- Executable entry point+-- ---------------------------------------------------------------------------++-- | The demo's command-line core. With no arguments it runs the pipeline live,+-- prints the trace tree and the final critique, and writes @trace.json@. With+-- @--replay PATH@ it loads the trace, replays the pipeline offline, and prints+-- the same critique plus @provider calls: 0@.+demoMain :: [String] -> IO ()+demoMain args = case args of+ [] -> liveMode+ ["--replay", path] -> replayMode path+ _ -> TIO.putStrLn "usage: shikumi-trace-demo [--replay PATH]"++liveMode :: IO ()+liveMode = do+ (critique, tree) <-+ runEff . runPrim . runTime . runTrace . runStubLLM . tracedLLM $ demoPipeline demoArticle+ TIO.putStr (renderTree tree)+ writeTraceFile "trace.json" tree+ TIO.putStrLn ("FINAL: " <> critique)++replayMode :: FilePath -> IO ()+replayMode path = do+ _ <- lookupEnv "SHIKUMI_OFFLINE" -- offline is structural: replay contacts no provider+ loaded <- readTraceFile path+ case loaded of+ Left err -> TIO.putStrLn ("trace load error: " <> err)+ Right tree -> do+ let idx = replayIndex tree+ (critique, _) <- runEff . runPrim . runTime . runTrace . runLLMReplay idx $ demoPipeline demoArticle+ TIO.putStrLn ("FINAL: " <> critique)+ TIO.putStrLn "provider calls: 0"++-- ---------------------------------------------------------------------------+-- helpers+-- ---------------------------------------------------------------------------++opts :: Options+opts = _Options++ctxFor :: Text -> Context+ctxFor t = _Context & #messages .~ V.singleton (user t)++mkResponse :: Text -> Response+mkResponse t =+ _Response+ & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))+ & #message . #usage . #inputTokens .~ 24+ & #message . #usage . #outputTokens .~ 9+ & #latencyMs .~ 7++textOf :: Response -> Text+textOf resp =+ mconcat [t | AssistantText (TextContent t) <- V.toList (resp ^. #message . #content)]++lastUserText :: Context -> Text+lastUserText c =+ mconcat+ [ t+ | UserMessage up <- V.toList (c ^. #messages),+ UserText (TextContent t) <- V.toList (up ^. #content)+ ]
+ src/Shikumi/Trace/Feedback.hs view
@@ -0,0 +1,78 @@+{-# LANGUAGE GADTs #-}+{-# LANGUAGE TypeFamilies #-}++-- | The per-node feedback channel (EP-16, M3).+--+-- A 'FeedbackLog' is a sibling of the trace (deliberately /not/ part of+-- 'Shikumi.Trace.TraceTree', whose serialized format is pinned and consumed by+-- replay): it maps a node's 'NodePath' to the textual critiques attached to it. A+-- metric or LM-judge writes a critique with 'attachFeedback' during evaluation; an+-- optimizer reads all critiques for a node with 'feedbackFor' during a proposal+-- step. GEPA (@docs/plans/22-gepa-reflective-optimizer.md@) is the headline+-- consumer.+module Shikumi.Trace.Feedback+ ( FeedbackLog (..),+ emptyFeedback,+ feedbackFor,+ Feedback (..),+ attachFeedback,+ runFeedback,+ )+where++import Data.Aeson (FromJSON (..), ToJSON (..))+import Data.Map.Strict (Map)+import Data.Map.Strict qualified as Map+import Data.Text (Text)+import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>))+import Effectful.Dispatch.Dynamic (interpret, send)+import Effectful.Prim (Prim)+import Effectful.Prim.IORef (modifyIORef', newIORef, readIORef)+import Shikumi.Trace.Node (NodePath)++-- | All feedback gathered during a run: critiques keyed by the node they target. A+-- node may accumulate several critiques (e.g. one per example), kept in attach+-- order.+newtype FeedbackLog = FeedbackLog (Map NodePath [Text])+ deriving stock (Eq, Show)++-- | Serialized as a list of @(NodePath, [Text])@ pairs, so 'NodePath' needs only+-- 'ToJSON'\/'FromJSON' (not a text 'ToJSONKey').+instance ToJSON FeedbackLog where+ toJSON (FeedbackLog m) = toJSON (Map.toList m)++instance FromJSON FeedbackLog where+ parseJSON v = FeedbackLog . Map.fromList <$> parseJSON v++-- | The empty log.+emptyFeedback :: FeedbackLog+emptyFeedback = FeedbackLog Map.empty++-- | All critiques attached to a node, in attach order (empty if none).+feedbackFor :: NodePath -> FeedbackLog -> [Text]+feedbackFor p (FeedbackLog m) = Map.findWithDefault [] p m++-- | The effect a metric or judge uses to attach feedback during a run.+data Feedback :: Effect where+ AttachFeedback :: NodePath -> Text -> Feedback m ()++type instance DispatchOf Feedback = 'Dynamic++-- | Attach a short textual critique to a node.+attachFeedback :: (Feedback :> es) => NodePath -> Text -> Eff es ()+attachFeedback p t = send (AttachFeedback p t)++-- | Run the feedback effect, returning the result paired with the collected log.+-- Accumulates into an @IORef@ via 'Prim' (the same mechanism @runTrace@ uses),+-- appending each critique to its node's list.+runFeedback :: (Prim :> es) => Eff (Feedback : es) a -> Eff es (a, FeedbackLog)+runFeedback act = do+ ref <- newIORef Map.empty+ a <-+ interpret+ ( \_ -> \case+ AttachFeedback p t -> modifyIORef' ref (Map.insertWith (\new old -> old ++ new) p [t])+ )+ act+ m <- readIORef ref+ pure (a, FeedbackLog m)
+ src/Shikumi/Trace/Internal/Spike.hs view
@@ -0,0 +1,142 @@+{-# LANGUAGE GADTs #-}++-- | M0 de-risking spike for EP-7+-- (@docs/plans/7-hierarchical-tracing-observability-and-replay.md@).+--+-- The one genuinely novel piece of the tracing plan is /how the span hierarchy is+-- formed/, since neither baikai nor EP-1's @LLM@ effect carries a parent\/child+-- relationship. This spike proves the mechanism the production effect (M1) is+-- built on:+--+-- * a manually-maintained __stack of span ids__ (an @IORef [SpanId]@), and+-- * an __interpose__ over EP-1's @LLM@ effect that, on each 'Shikumi.LLM.complete',+-- reads the id currently on top of the stack and tags the call with it before+-- delegating to the real handler.+--+-- This is the same @interpose@ seam EP-6's @cachedLLM@ already uses, so the risk+-- is low; the spike confirms that the id on top of the stack /at the moment the+-- call is sent/ is the enclosing span, which is exactly what M1 needs.+--+-- Decision (EP-7, 2026-06-08): capture LM calls by interposing on the @LLM@+-- effect rather than installing a baikai 'Baikai.Trace.Sink.TraceSink'. EP-1's+-- interpreters expose no sink parameter, and @LLM.complete@ returns the full+-- baikai 'Baikai.Response' (latency, usage, cost, tool blocks) — strictly more+-- than baikai's flat @TraceEvent@. See the plan's Decision Log.+module Shikumi.Trace.Internal.Spike+ ( SpanId (..),+ runSpike,+ )+where++import Baikai+ ( AssistantContent (..),+ Context,+ Response,+ TextContent (..),+ _Context,+ _Model,+ _Options,+ _Response,+ _TextContent,+ )+import Control.Lens ((&), (.~), (^.))+import Data.Generics.Labels ()+import Data.IORef (IORef, modifyIORef', newIORef, readIORef)+import Data.Text (Text)+import Data.Vector qualified as V+import Effectful (Eff, IOE, liftIO, runEff, (:>))+import Effectful.Dispatch.Dynamic (interpose, interpret)+import Shikumi.LLM (LLM (..), complete)++-- | A span identifier. The production effect (M1) reuses this newtype.+newtype SpanId = SpanId Text+ deriving stock (Eq, Ord, Show)++-- | Run two fake LM calls nested under two distinct span ids and return, for+-- each captured call, the @(enclosing span id, response text)@ pair. If the+-- mechanism works, the first call is tagged with @"draft"@ and the second with+-- @"critique"@ — the ids that were on top of the stack when each call ran.+runSpike :: IO [(SpanId, Text)]+runSpike = runEff $ do+ stack <- liftIO (newIORef [])+ captured <- liftIO (newIORef [])+ runScripted [stubResponse "first draft", stubResponse "a sharp critique"]+ . traceSpike stack captured+ $ do+ _ <- withSpike stack (SpanId "draft") oneCall+ _ <- withSpike stack (SpanId "critique") oneCall+ pure ()+ liftIO (readIORef captured)+ where+ oneCall = complete _Model neutralCtx _Options++neutralCtx :: Context+neutralCtx = _Context++-- ---------------------------------------------------------------------------+-- The mechanism under test+-- ---------------------------------------------------------------------------++-- | Push @sid@ on the stack for the dynamic extent of @act@, then pop it. The+-- enclosing span is whatever is on top while @act@ runs.+withSpike :: (IOE :> es) => IORef [SpanId] -> SpanId -> Eff es a -> Eff es a+withSpike stack sid act = do+ liftIO (modifyIORef' stack (sid :))+ r <- act+ liftIO (modifyIORef' stack (drop 1))+ pure r++-- | Interpose on @LLM@: on each 'Complete', read the top-of-stack span id,+-- delegate to the underlying handler, and record @(top, responseText)@.+traceSpike ::+ (IOE :> es, LLM :> es) =>+ IORef [SpanId] ->+ IORef [(SpanId, Text)] ->+ Eff es a ->+ Eff es a+traceSpike stack captured = interpose $ \_ -> \case+ Complete m c o -> do+ top <- topOf <$> liftIO (readIORef stack)+ resp <- complete m c o+ liftIO (modifyIORef' captured (++ [(top, responseText resp)]))+ pure resp+ Stream m c o -> stream' m c o+ where+ topOf [] = SpanId "<root>"+ topOf (x : _) = x+ stream' m c o = complete m c o >> pure []++-- ---------------------------------------------------------------------------+-- A network-free scripted @LLM@ interpreter (pops canned responses in order)+-- ---------------------------------------------------------------------------++runScripted :: (IOE :> es) => [Response] -> Eff (LLM : es) a -> Eff es a+runScripted rs0 act = do+ ref <- liftIO (newIORef rs0)+ interpret+ ( \_ -> \case+ Complete {} -> liftIO (pop ref)+ Stream {} -> pure []+ )+ act+ where+ pop :: IORef [Response] -> IO Response+ pop ref = do+ rs <- readIORef ref+ case rs of+ (x : xs) -> modifyIORef' ref (const xs) >> pure x+ [] -> pure (stubResponse "")++-- ---------------------------------------------------------------------------+-- Tiny helpers+-- ---------------------------------------------------------------------------++-- | A 'Response' carrying @t@ as its single assistant text block.+stubResponse :: Text -> Response+stubResponse t =+ _Response & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))++-- | Concatenate the assistant text blocks of a response.+responseText :: Response -> Text+responseText resp =+ mconcat [t | AssistantText (TextContent t) <- V.toList (resp ^. #message . #content)]
+ src/Shikumi/Trace/Node.hs view
@@ -0,0 +1,126 @@+{-# LANGUAGE ScopedTypeVariables #-}++-- | Node identity for the trace (EP-16, M1).+--+-- A 'NodePath' names a 'Shikumi.Program.Program' node's structural position as the+-- list of branch steps from the program root. 'programNodePaths' enumerates the+-- path of every @Predict@ node in /exactly/ the left-to-right depth-first order+-- 'Shikumi.Program.foldParams' yields their @Params@ — so @programNodePaths p !! n@+-- is the path to the same node @mapParamsAt n@ edits and @foldParams p !! n@+-- parameterizes. That agreement is the parent MasterPlan's integration point #3,+-- and it is structural (the enumeration reuses the identical depth-first descent),+-- not a coincidence two functions must independently maintain.+module Shikumi.Trace.Node+ ( NodeStep (..),+ NodePath (..),+ programNodePaths,+ renderNodePath,+ nodeFields,+ )+where++import Data.Aeson (FromJSON, ToJSON)+import Data.Text (Text)+import Data.Text qualified as T+import GHC.Generics (Generic)+import Shikumi.Program+ ( NodeFields,+ Program+ ( Compose,+ Embed,+ Ensemble,+ FMap,+ MajorityVote,+ Map,+ Parallel,+ Predict,+ Retry,+ RetryWhen,+ Validate+ ),+ nodeFieldsIndexed,+ )++-- | One structural step from a parent node to a child, naming which branch was+-- taken. The labels mirror the 'Shikumi.Program.Program' constructors so a path is+-- human-readable and shape-stable: two programs of the same shape yield identical+-- paths.+data NodeStep+ = -- | into the left side of @Compose@+ StepComposeL+ | -- | into the right side of @Compose@+ StepComposeR+ | -- | through an @FMap@+ StepFMap+ | -- | through a @Map@+ StepMap+ | -- | into the left side of @Parallel@+ StepParallelL+ | -- | into the right side of @Parallel@+ StepParallelR+ | -- | through a @Retry@+ StepRetry+ | -- | through a @RetryWhen@+ StepRetryWhen+ | -- | through a @Validate@+ StepValidate+ | -- | through a @MajorityVote@+ StepMajorityVote+ | -- | into the i-th member of an @Ensemble@ (0-based)+ StepEnsemble !Int+ deriving stock (Eq, Ord, Show, Generic)+ deriving anyclass (ToJSON, FromJSON)++-- | The structural position of a node within a 'Shikumi.Program.Program', as the+-- list of steps from the program root to that node, outermost first. A bare+-- root @Predict@ has the empty path.+newtype NodePath = NodePath [NodeStep]+ deriving stock (Eq, Ord, Show, Generic)+ deriving newtype (ToJSON, FromJSON)++-- | Enumerate the 'NodePath' of every @Predict@ node, in the same left-to-right+-- depth-first order as 'Shikumi.Program.foldParams'. Reuses that identical descent+-- so length and node order agree by construction.+programNodePaths :: Program i o -> [NodePath]+programNodePaths = go []+ where+ -- @prefix@ is the reversed list of steps taken so far; we reverse on emit.+ go :: forall x y. [NodeStep] -> Program x y -> [NodePath]+ go prefix (Predict _ _) = [NodePath (reverse prefix)]+ go prefix (Compose a b) = go (StepComposeL : prefix) a ++ go (StepComposeR : prefix) b+ go prefix (FMap _ p) = go (StepFMap : prefix) p+ go prefix (Map _ p) = go (StepMap : prefix) p+ go prefix (Parallel a b) = go (StepParallelL : prefix) a ++ go (StepParallelR : prefix) b+ go prefix (Retry _ p) = go (StepRetry : prefix) p+ go prefix (RetryWhen _ _ p) = go (StepRetryWhen : prefix) p+ go prefix (Validate _ p) = go (StepValidate : prefix) p+ go prefix (MajorityVote _ _ p) = go (StepMajorityVote : prefix) p+ go prefix (Ensemble ps _) = concat (zipWith (\i p -> go (StepEnsemble i : prefix) p) [0 ..] ps)+ go _ (Embed _) = []++-- | Render a 'NodePath' to a short, stable string (e.g. @compose.0/predict@-style+-- slash-joined steps), suitable as a trace/OTel attribute value. The empty path+-- renders as @\"root\"@.+renderNodePath :: NodePath -> Text+renderNodePath (NodePath []) = "root"+renderNodePath (NodePath steps) = T.intercalate "/" (map stepLabel steps)+ where+ stepLabel = \case+ StepComposeL -> "compose.l"+ StepComposeR -> "compose.r"+ StepFMap -> "fmap"+ StepMap -> "map"+ StepParallelL -> "parallel.l"+ StepParallelR -> "parallel.r"+ StepRetry -> "retry"+ StepRetryWhen -> "retryWhen"+ StepValidate -> "validate"+ StepMajorityVote -> "majorityVote"+ StepEnsemble i -> "ensemble." <> T.pack (show i)++-- | Associate each @Predict@ node's 'NodePath' with its input/output field names.+-- A consumer with a 'NodePath' from a trace span zips against this to map+-- path → field metadata; one with a @foldParams@ index uses+-- 'Shikumi.Program.nodeFieldsIndexed' directly.+nodeFields :: Program i o -> [(NodePath, NodeFields)]+nodeFields p = zip (programNodePaths p) (nodeFieldsIndexed p)
+ src/Shikumi/Trace/Program.hs view
@@ -0,0 +1,186 @@+{-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilies #-}++-- | Node-correlated program execution for the trace (EP-16, M2).+--+-- 'runProgramTraced' is an /additive/ entry point: it runs a+-- 'Shikumi.Program.Program' exactly like 'Shikumi.Program.runProgram' — reusing+-- its render\/parse\/retry\/vote semantics by delegating each @Predict@ leaf to+-- @runProgram@ — but it also opens a trace span per node and threads the active+-- node's 'NodePath' so that each model-call span is tagged with the structural+-- position of the node that issued it. 'Shikumi.Program.runProgram' /+-- 'Shikumi.Program.runProgramConc' are untouched (MasterPlan integration point #4).+--+-- The correlation is threaded by a tiny reader-style 'CurrentNode' effect held in+-- an @IORef@ cell: 'runProgramTraced' wraps each @Predict@ in 'localNode' as it+-- descends (with the same step prefix 'Shikumi.Trace.Node.programNodePaths' uses),+-- and the node-aware capture 'tracedNodeLLM' reads it ('askNode') to stamp the+-- span. The capture and the tag happen inside one 'withSpan', so there is no+-- cross-interpreter ordering subtlety (see the Decision Log of EP-16): use+-- 'tracedNodeLLM' /instead of/ 'Shikumi.Trace.tracedLLM' for a traced-program run.+--+-- Stack shape (note 'runCurrentNode' is outer of 'tracedNodeLLM' so both share the+-- same node cell):+--+-- @+-- runEff . runPrim . runTime . runTrace . runCurrentNode+-- . runKeyedLLM responder -- base LLM interpreter (stub or real)+-- . tracedNodeLLM -- opens the LlmCallSpan, fills attrs, tags nodePath+-- $ runProgramTraced program input+-- @+module Shikumi.Trace.Program+ ( -- * The current-node effect+ CurrentNode (..),+ askNode,+ localNode,+ runCurrentNode,++ -- * Node-aware capture and execution+ tracedNodeLLM,+ runProgramTraced,+ )+where++import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>))+import Effectful.Dispatch.Dynamic (interpose, interpret, localSeqUnlift, send)+import Effectful.Error.Static (Error)+import Effectful.Exception (bracket_)+import Effectful.Prim (Prim)+import Effectful.Prim.IORef (newIORef, readIORef, writeIORef)+import Shikumi.Error (ShikumiError)+import Shikumi.LLM (LLM (..), complete, stream)+import Shikumi.Program+ ( Program+ ( Compose,+ Embed,+ Ensemble,+ FMap,+ MajorityVote,+ Map,+ Parallel,+ Predict,+ Retry,+ RetryWhen,+ Validate+ ),+ acceptOrReject,+ modal,+ retryWith,+ runProgram,+ sampleTemps,+ withSampleTemp,+ )+import Shikumi.Trace+ ( SpanAttrs (..),+ SpanKind (CombinatorSpan, LlmCallSpan, ModuleSpan),+ Trace,+ annotateSpan,+ llmAttrs,+ llmLabel,+ withSpan,+ )+import Shikumi.Trace.Node (NodePath (..), NodeStep (..))++-- ---------------------------------------------------------------------------+-- The current-node effect+-- ---------------------------------------------------------------------------++-- | Carries the 'NodePath' of the node currently executing. 'localNode' sets it+-- for the duration of an inner action; 'askNode' reads it. A dedicated effect+-- (rather than @Effectful.Reader@) keeps the value dynamically scoped as execution+-- descends the program tree without the caller threading it, and keeps+-- 'Shikumi.Program.runProgram'\'s row untouched.+data CurrentNode :: Effect where+ AskNode :: CurrentNode m (Maybe NodePath)+ LocalNode :: NodePath -> m a -> CurrentNode m a++type instance DispatchOf CurrentNode = 'Dynamic++-- | The node currently executing, if any.+askNode :: (CurrentNode :> es) => Eff es (Maybe NodePath)+askNode = send AskNode++-- | Run an action with the given node marked as current.+localNode :: (CurrentNode :> es) => NodePath -> Eff es a -> Eff es a+localNode p act = send (LocalNode p act)++-- | Discharge 'CurrentNode' with a mutable cell (via 'Prim'). Because the cell is+-- shared, an 'askNode' issued from a lower interpose ('tracedNodeLLM') sees the+-- value a higher 'localNode' set, which is exactly how a model call is correlated+-- to its issuing node.+runCurrentNode :: (Prim :> es) => Eff (CurrentNode : es) a -> Eff es a+runCurrentNode act = do+ ref <- newIORef Nothing+ interpret+ ( \env -> \case+ AskNode -> readIORef ref+ LocalNode p inner -> do+ old <- readIORef ref+ bracket_+ (writeIORef ref (Just p))+ (writeIORef ref old)+ (localSeqUnlift env (\unlift -> unlift inner))+ )+ act++-- ---------------------------------------------------------------------------+-- Node-aware capture+-- ---------------------------------------------------------------------------++-- | Like 'Shikumi.Trace.tracedLLM' but also stamps the active 'NodePath' onto each+-- model-call span. The capture (model\/prompt\/response\/cost) and the node tag are+-- written inside the /same/ 'withSpan', so the tag always lands on the LM-call span+-- — no dependence on interpose ordering relative to a separate capture layer.+tracedNodeLLM :: (Trace :> es, CurrentNode :> es, LLM :> es) => Eff es a -> Eff es a+tracedNodeLLM = interpose $ \_ -> \case+ Complete m c o -> withSpan LlmCallSpan (llmLabel m) $ do+ resp <- complete m c o+ mp <- askNode+ annotateSpan (\_ -> (llmAttrs m c o resp) {nodePath = mp})+ pure resp+ Stream m c o -> withSpan LlmCallSpan (llmLabel m) (stream m c o)++-- ---------------------------------------------------------------------------+-- Node-correlated execution+-- ---------------------------------------------------------------------------++-- | Run a program like 'Shikumi.Program.runProgram', opening a span per node and+-- tagging each model-call span with the issuing node's 'NodePath'. The step prefix+-- accumulated as it descends is the same one 'Shikumi.Trace.Node.programNodePaths'+-- builds, so a @Predict@ leaf's path here equals the path that enumeration assigns+-- it. Each @Predict@ leaf delegates to @runProgram@ (reusing its exact semantics);+-- the combinators are mirrored with spans and the same control flow.+runProgramTraced ::+ forall i o es.+ (LLM :> es, Trace :> es, CurrentNode :> es, Error ShikumiError :> es) =>+ Program i o ->+ i ->+ Eff es o+runProgramTraced = go []+ where+ go :: forall x y. [NodeStep] -> Program x y -> x -> Eff es y+ go prefix node@(Predict _ _) i =+ withSpan ModuleSpan "Predict" (localNode (NodePath (reverse prefix)) (runProgram node i))+ go prefix (Compose f g) i =+ withSpan CombinatorSpan "Compose" (go (StepComposeL : prefix) f i >>= go (StepComposeR : prefix) g)+ go prefix (FMap k p) i =+ withSpan CombinatorSpan "FMap" (k <$> go (StepFMap : prefix) p i)+ go prefix (Map _ p) xs =+ withSpan CombinatorSpan "Map" (traverse (go (StepMap : prefix) p) xs)+ go prefix (Parallel a b) i =+ withSpan CombinatorSpan "Parallel" ((,) <$> go (StepParallelL : prefix) a i <*> go (StepParallelR : prefix) b i)+ go prefix (Retry n p) i =+ withSpan CombinatorSpan "Retry" (retryWith (go (StepRetry : prefix)) (const True) n p i)+ go prefix (RetryWhen ok n p) i =+ withSpan CombinatorSpan "RetryWhen" (retryWith (go (StepRetryWhen : prefix)) ok n p i)+ go prefix (Validate v p) i =+ withSpan CombinatorSpan "Validate" (go (StepValidate : prefix) p i >>= acceptOrReject v)+ go prefix (MajorityVote k sched p) i =+ withSpan CombinatorSpan "MajorityVote" $+ modal <$> traverse (\mt -> withSampleTemp mt (go (StepMajorityVote : prefix) p i)) (sampleTemps (max 1 k) sched)+ go prefix (Ensemble ps reduce) i =+ withSpan CombinatorSpan "Ensemble" $+ reduce <$> sequence [go (StepEnsemble idx : prefix) p i | (idx, p) <- zip [0 ..] ps]+ go _ (Embed f) i =+ withSpan CombinatorSpan "Embed" (f i)
+ src/Shikumi/Trace/Replay.hs view
@@ -0,0 +1,95 @@+{-# LANGUAGE GADTs #-}++-- | Deterministic offline replay of the @LLM@ effect (EP-7, M3).+--+-- 'runLLMReplay' is a drop-in /alternative interpreter/ of EP-1's @LLM@ effect:+-- where the real interpreter dispatches to a provider, this one answers every+-- 'Shikumi.LLM.complete' from a recorded trace. It computes the EP-6+-- content-addressed 'Shikumi.Cache.Key.cacheKey' of the request, looks it up in a+-- replay index (built by 'Shikumi.Trace.Store.replayIndex'), and returns the+-- recorded 'Baikai.Response' — decoded via the 'Shikumi.Trace.ResponseJSON'+-- instances. Because it replaces the @LLM@ effect at the /same boundary/, the rest+-- of the program (modules, combinators, decoding) runs exactly as in live mode, so+-- the typed outputs are identical; only the leaf LM calls are redirected.+--+-- Divergence is __fail-closed and loud__: a request whose key is not in the trace+-- raises a typed 'ReplayDivergence' carrying the key, the model id, and a redacted+-- prompt summary. It never falls through to the network and never fabricates a+-- response. (There is no registry in this interpreter at all, so "zero provider+-- calls" is structural, not merely policy.)+module Shikumi.Trace.Replay+ ( ReplayDivergence (..),+ runLLMReplay,+ )+where++import Baikai (Context, Model, Response)+import Control.Lens ((^.))+import Data.Aeson (Result (Error, Success), Value, fromJSON)+import Data.Generics.Labels ()+import Data.Map.Strict (Map)+import Data.Map.Strict qualified as Map+import Data.Text (Text)+import Data.Text qualified as T+import Effectful (Eff)+import Effectful.Dispatch.Dynamic (interpret)+import Effectful.Exception (Exception, throwIO)+import GHC.Generics (Generic)+import Shikumi.Cache.Key (CacheKey (..), cacheKey)+import Shikumi.LLM (LLM (..))+import Shikumi.Trace.ResponseJSON ()++-- | Raised when a replayed program issues an LM call that the recorded trace+-- never saw.+data ReplayDivergence = ReplayDivergence+ { -- | the content-addressed key that was absent from the trace+ divergedKey :: !CacheKey,+ -- | the model the diverging call targeted+ divergedModel :: !Text,+ -- | a redacted one-line summary of the diverging request+ promptSummary :: !Text+ }+ deriving stock (Eq, Show, Generic)+ deriving anyclass (Exception)++-- | Interpret the @LLM@ effect by lookup in a replay index instead of calling a+-- provider. A hit returns the recorded response; a miss raises 'ReplayDivergence'.+-- Streaming completions are not replayable and raise a divergence naming the key.+--+-- No @IOE@ is required: the lookup and decode are pure, and 'ReplayDivergence' is+-- raised with @effectful@'s pure-in-@Eff@ 'throwIO'. The contract the plan sketched+-- as @(IOE :> es)@ is therefore satisfied with a strictly weaker constraint.+runLLMReplay :: Map CacheKey Value -> Eff (LLM : es) a -> Eff es a+runLLMReplay idx = interpret $ \_ -> \case+ Complete m c o -> do+ let key = cacheKey m c o+ case Map.lookup key idx of+ Nothing -> throwIO (divergence key m c)+ Just v -> case fromJSON v of+ Success r -> pure (r :: Response)+ Error e ->+ throwIO+ (divergence key m c) {promptSummary = "recorded response failed to decode: " <> T.pack e}+ Stream m c o ->+ throwIO+ (divergence (cacheKey m c o) m c) {promptSummary = "replay does not support streaming completions"}++-- | Build a 'ReplayDivergence' for a request.+divergence :: CacheKey -> Model -> Context -> ReplayDivergence+divergence key m c =+ ReplayDivergence+ { divergedKey = key,+ divergedModel = m ^. #modelId,+ promptSummary = summarizeRequest m c+ }++-- | A redacted, one-line request summary for divergence errors: the routing+-- identity and the message count (no message bodies — those may be sensitive).+summarizeRequest :: Model -> Context -> Text+summarizeRequest m c =+ (m ^. #provider)+ <> "/"+ <> (m ^. #modelId)+ <> " ("+ <> T.pack (show (length (c ^. #messages)))+ <> " messages)"
+ src/Shikumi/Trace/ResponseJSON.hs view
@@ -0,0 +1,17 @@+-- | A faithful JSON round-trip for baikai's 'Baikai.Response.Response' graph.+--+-- This module exists for backward compatibility: it re-exports the baikai+-- 'Response'-graph orphan instances ('ToJSON'/'FromJSON' for @Response@,+-- @AssistantPayload@, @Usage@, @Cost@, @CostBreakdown@) from their single home,+-- "Shikumi.Cache.ResponseJSON" in @shikumi-cache@ (EP-6 owns the cache key that+-- @shikumi-trace@ already depends on, so the dependency direction is natural).+--+-- The instances originally lived here (EP-7). They moved down to @shikumi-cache@+-- when EP-6's persistent backends — which need the same round-trip to store a+-- response as JSON — landed, so that a module importing both the cache and the+-- trace package never sees a duplicate orphan instance. Importing this module+-- (@import Shikumi.Trace.ResponseJSON ()@) still brings the instances into scope+-- exactly as before.+module Shikumi.Trace.ResponseJSON () where++import Shikumi.Cache.ResponseJSON ()
+ src/Shikumi/Trace/Store.hs view
@@ -0,0 +1,91 @@+-- | Persisting a 'TraceTree' to a stable on-disk JSON format, and deriving the+-- replay index from it (EP-7, M2).+--+-- A trace file is a single JSON document — a 'TraceFile' carrying a+-- @formatVersion@ integer and the whole 'TraceTree'. Persisting the tree as a+-- tree (rather than a flat event log) keeps replay a simple key lookup and keeps+-- the file inspectable with @jq@ (e.g. @jq '.tree.spans | length'@).+--+-- 'writeTraceFile' is atomic (write to a sibling @.tmp@ then rename), so a crash+-- mid-write never leaves a half-written file. 'readTraceFile' returns 'Left' on a+-- parse error or a format-version mismatch — replay refuses to start on a file it+-- does not understand rather than producing wrong outputs.+--+-- 'replayIndex' projects the tree into the @Map CacheKey Value@ the replay+-- interpreter ('Shikumi.Trace.Replay') consults: each LM-call span's EP-6+-- 'Shikumi.Cache.Key.CacheKey' mapped to its recorded response JSON.+module Shikumi.Trace.Store+ ( TraceFile (..),+ currentFormatVersion,+ writeTraceFile,+ readTraceFile,+ replayIndex,+ )+where++import Control.Lens ((^.))+import Data.Aeson (FromJSON, ToJSON, Value, eitherDecode, encode)+import Data.ByteString.Lazy qualified as BL+import Data.Generics.Labels ()+import Data.Map.Strict (Map)+import Data.Map.Strict qualified as Map+import Data.Text (Text)+import Data.Text qualified as T+import GHC.Generics (Generic)+import Shikumi.Cache.Key (CacheKey (..))+import Shikumi.Trace (SpanKind (LlmCallSpan), TraceTree (..))+import System.Directory (renameFile)++-- | The on-disk wrapper: a format version plus the tree. Bump+-- 'currentFormatVersion' on any breaking schema change so old files are rejected+-- loudly rather than mis-decoded.+data TraceFile = TraceFile+ { formatVersion :: !Int,+ tree :: !TraceTree+ }+ deriving stock (Eq, Show, Generic)+ deriving anyclass (ToJSON, FromJSON)++-- | The trace-file schema version this build reads and writes. Bumped 1→2 in+-- EP-16 for the additive optional @SpanAttrs.nodePath@ field.+currentFormatVersion :: Int+currentFormatVersion = 2++-- | Write a tree to @path@ atomically: encode to a sibling @.tmp@ file, then+-- 'renameFile' it into place (an atomic rename on the same filesystem).+writeTraceFile :: FilePath -> TraceTree -> IO ()+writeTraceFile path t = do+ let tmp = path <> ".tmp"+ BL.writeFile tmp (encode (TraceFile currentFormatVersion t))+ renameFile tmp path++-- | Read a tree from @path@. 'Left' on a JSON parse error or a @formatVersion@+-- this build does not understand.+readTraceFile :: FilePath -> IO (Either Text TraceTree)+readTraceFile path = do+ bs <- BL.readFile path+ pure $ case eitherDecode bs of+ Left e -> Left ("trace parse error: " <> T.pack e)+ Right tf+ | formatVersion tf /= currentFormatVersion ->+ Left+ ( "unsupported trace formatVersion "+ <> T.pack (show (formatVersion tf))+ <> " (expected "+ <> T.pack (show currentFormatVersion)+ <> ")"+ )+ | otherwise -> Right (tree tf)++-- | The replay index: every LM-call span's content-addressed key mapped to its+-- recorded response JSON. A span missing either its @cacheKey@ or its @response@+-- contributes nothing.+replayIndex :: TraceTree -> Map CacheKey Value+replayIndex t =+ Map.fromList+ [ (CacheKey ck, v)+ | s <- Map.elems (spans t),+ (s ^. #kind) == LlmCallSpan,+ Just ck <- [s ^. #attrs . #cacheKey],+ Just v <- [s ^. #attrs . #response]+ ]
+ test/Main.hs view
@@ -0,0 +1,533 @@+-- | EP-7 acceptance test suite. Milestones are selected with tasty's @-p@:+-- @spike@ (M0), @tree@ (M1), @store@ (M2), @replay@ (M3), @e2e@ (M5).+module Main (main) where++import Baikai (Context, Model, Options, Response, user, _Context, _Model, _Options)+import Control.Exception (try)+import Control.Lens ((&), (.~))+import Data.Aeson (Value (Object), decode, encode, object, (.=))+import Data.ByteString.Lazy qualified as BL+import Data.Generics.Labels ()+import Data.IORef (newIORef, readIORef, writeIORef)+import Data.List (sortOn)+import Data.Map.Strict qualified as Map+import Data.Text (Text)+import Data.Text qualified as T+import Data.Time.Clock (UTCTime, addUTCTime)+import Data.Vector qualified as V+import Effectful (Eff, runEff, type (:>))+import Effectful.Error.Static (runErrorNoCallStack)+import Effectful.Prim (runPrim)+import GHC.Generics (Generic)+import Shikumi.Adapter (ToPrompt)+import Shikumi.Cache.Key (CacheKey (..))+import Shikumi.Cache.Key qualified as Key+import Shikumi.Combinator (majorityVote, majorityVoteBy, parallel2, (>>>))+import Shikumi.Effect.Time (runTime)+import Shikumi.Error (ShikumiError)+import Shikumi.LLM (LLM, complete)+import Shikumi.Module (predict)+import Shikumi.Program+ ( NodeFields (..),+ Params (..),+ Program,+ TempSchedule (TempFixed),+ foldParams,+ mapParamsAt,+ )+import Shikumi.Schema (FromModel, ToSchema, Validatable)+import Shikumi.Signature (Signature, mkSignature)+import Shikumi.Trace+ ( Span (..),+ SpanAttrs (..),+ SpanId (..),+ SpanKind (..),+ TraceTree (..),+ childrenOf,+ emptyAttrs,+ renderTree,+ runTrace,+ tracedLLM,+ withSpan,+ )+import Shikumi.Trace.Demo (demoArticle, demoPipeline, demoResponder)+import Shikumi.Trace.Feedback+ ( FeedbackLog,+ attachFeedback,+ feedbackFor,+ runFeedback,+ )+import Shikumi.Trace.Internal.Spike qualified as Spike+import Shikumi.Trace.Node+ ( NodePath (..),+ NodeStep (..),+ nodeFields,+ programNodePaths,+ )+import Shikumi.Trace.Program+ ( runCurrentNode,+ runProgramTraced,+ tracedNodeLLM,+ )+import Shikumi.Trace.Replay (ReplayDivergence (..), runLLMReplay)+import Shikumi.Trace.Store+ ( TraceFile (..),+ currentFormatVersion,+ readTraceFile,+ replayIndex,+ writeTraceFile,+ )+import System.IO.Temp (withSystemTempDirectory)+import Test.Tasty (TestTree, defaultMain, testGroup)+import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=))+import Test.Tasty.QuickCheck (Gen, choose, elements, forAll, testProperty, vectorOf, (===))+import TraceFixtures+ ( ctxFor,+ lastUserText,+ mkResponse,+ optsFor,+ responseText,+ runFixedLLM,+ runKeyedCountingLLM,+ stubModel,+ )++main :: IO ()+main =+ defaultMain $+ testGroup+ "shikumi-trace"+ [spikeTests, treeTests, storeTests, replayTests, e2eTests, nodeTests, correlateTests, feedbackTests]++-- ---------------------------------------------------------------------------+-- M0+-- ---------------------------------------------------------------------------++-- | M0: the interpose-plus-span-stack mechanism tags each LM call with the+-- enclosing span that was on top of the stack when the call ran.+spikeTests :: TestTree+spikeTests =+ testGroup+ "spike"+ [ testCase "each captured LM call is tagged with its enclosing span id" $ do+ captured <- Spike.runSpike+ map fst captured @?= [Spike.SpanId "draft", Spike.SpanId "critique"]+ map snd captured @?= ["first draft", "a sharp critique"]+ ]++-- ---------------------------------------------------------------------------+-- M1+-- ---------------------------------------------------------------------------++-- | M1: a two-@withSpan@ nest containing two stubbed LM calls produces a tree+-- with one root, two module children, each holding exactly one LM-call leaf;+-- @renderTree@ shows the model lines indented under their modules.+treeTests :: TestTree+treeTests =+ testGroup+ "tree"+ [ testCase "nested withSpans build a program -> module -> llm-call tree" $ do+ tree <- buildTree+ let roots = [s | s <- Map.elems (spans tree), parent s == Nothing]+ length roots @?= 1+ kind (spanAt tree (root tree)) @?= ProgramSpan+ let kids = childrenOf tree (root tree)+ length kids @?= 2+ map (kind . spanAt tree) kids @?= [ModuleSpan, ModuleSpan]+ let grandkids = concatMap (childrenOf tree) kids+ length grandkids @?= 2+ map (kind . spanAt tree) grandkids @?= [LlmCallSpan, LlmCallSpan],+ testCase "renderTree shows the model lines indented under their modules" $ do+ tree <- buildTree+ let out = renderTree tree+ assertBool "two llm-call model lines present" (T.count "stub/stub-model" out == 2)+ assertBool "llm-call lines are indented under modules" (T.isInfixOf " llm-call stub/stub-model" out)+ assertBool "module lines present" (T.count "module" out == 2)+ ]++-- ---------------------------------------------------------------------------+-- M2+-- ---------------------------------------------------------------------------++-- | M2: the trace tree serializes to a stable, versioned JSON document and reads+-- back; the replay index projects LM-call cacheKeys to responses; and the cache+-- key reproduces EP-6's pinned digest (integration point #7).+storeTests :: TestTree+storeTests =+ testGroup+ "store"+ [ testProperty "trace tree survives a JSON round-trip" $+ forAll genTree $ \t ->+ decode (encode (TraceFile currentFormatVersion t)) === Just (TraceFile currentFormatVersion t),+ testCase "writeTraceFile/readTraceFile round-trips on disk" $+ withSystemTempDirectory "shikumi-trace" $ \dir -> do+ tree <- buildTree+ let p = dir <> "/trace.json"+ writeTraceFile p tree+ res <- readTraceFile p+ res @?= Right tree,+ testCase "reading a file with a foreign formatVersion is Left" $+ withSystemTempDirectory "shikumi-trace" $ \dir -> do+ tree <- buildTree+ let p = dir <> "/bad.json"+ BL.writeFile p (encode (TraceFile 999 tree))+ res <- readTraceFile p+ case res of+ Left msg -> assertBool "the error names the offending version" (T.isInfixOf "999" msg)+ Right _ -> assertFailure "expected Left on a foreign formatVersion",+ testCase "replayIndex maps each llm-call cacheKey to its response" $ do+ tree <- buildTree+ let idx = replayIndex tree+ -- two LM-call leaves with distinct requests => two distinct keys.+ Map.size idx @?= 2+ assertBool "every indexed value is a JSON object (a recorded response)" (all isObject (Map.elems idx)),+ testCase "cacheKey reproduces EP-6's pinned digest (integration point #7)" $ do+ let CacheKey hex = Key.cacheKey fixModel fixCtx fixOpts+ hex @?= pinnedKey+ ]++-- ---------------------------------------------------------------------------+-- M3+-- ---------------------------------------------------------------------------++-- | M3: a two-stage pipeline captured live, persisted, and replayed offline+-- produces byte-identical outputs while making zero provider calls; a mutated+-- request raises a 'ReplayDivergence' naming the missing key.+replayTests :: TestTree+replayTests =+ testGroup+ "replay"+ [ testCase "replay reproduces the live outputs with zero provider calls" $+ withSystemTempDirectory "shikumi-trace" $ \dir -> do+ calls <- newIORef (0 :: Int)+ -- (1) live: run the pipeline, count provider calls, capture the tree.+ (live, tree) <-+ runEff . runPrim . runTime . runTrace . runKeyedCountingLLM calls responder . tracedLLM $+ twoStage "the article"+ liveCalls <- readIORef calls+ liveCalls @?= 2+ -- (2)+(3) persist and reload, build the replay index.+ let p = dir <> "/trace.json"+ writeTraceFile p tree+ Right tree' <- readTraceFile p+ let idx = replayIndex tree'+ -- (4) replay the SAME pipeline; the counting provider must not be hit.+ writeIORef calls 0+ replayed <- runEff . runLLMReplay idx $ twoStage "the article"+ replayCalls <- readIORef calls+ -- (5)+(6)+(7) identical outputs, zero provider calls.+ replayed @?= live+ replayCalls @?= 0,+ testCase "a request absent from the trace raises ReplayDivergence" $+ withSystemTempDirectory "shikumi-trace" $ \dir -> do+ calls <- newIORef (0 :: Int)+ (_, tree) <-+ runEff . runPrim . runTime . runTrace . runKeyedCountingLLM calls responder . tracedLLM $+ twoStage "the article"+ let p = dir <> "/trace.json"+ writeTraceFile p tree+ Right tree' <- readTraceFile p+ let idx = replayIndex tree'+ -- A different article => a first-stage key that was never recorded.+ res <- try (runEff . runLLMReplay idx $ twoStage "a different article")+ case res of+ Left (ReplayDivergence {divergedKey = CacheKey k}) ->+ assertBool "the divergence names a non-empty key" (not (T.null k))+ Right _ -> assertFailure "expected a ReplayDivergence on an unrecorded request"+ ]++-- ---------------------------------------------------------------------------+-- M5 (e2e)+-- ---------------------------------------------------------------------------++-- | M5: the demo's two-stage pipeline, captured live then replayed offline, prints+-- a byte-identical final critique while making zero provider calls — the headline+-- end-to-end story, run in-process.+e2eTests :: TestTree+e2eTests =+ testGroup+ "e2e"+ [ testCase "live capture then offline replay: identical FINAL, zero provider calls" $+ withSystemTempDirectory "shikumi-trace" $ \dir -> do+ calls <- newIORef (0 :: Int)+ -- live run of the demo pipeline; capture the tree, count provider calls.+ (liveFinal, tree) <-+ runEff . runPrim . runTime . runTrace . runKeyedCountingLLM calls demoResponder . tracedLLM $+ demoPipeline demoArticle+ liveCalls <- readIORef calls+ liveCalls @?= 2+ -- persist, reload, replay offline.+ let p = dir <> "/trace.json"+ writeTraceFile p tree+ Right tree' <- readTraceFile p+ let idx = replayIndex tree'+ writeIORef calls 0+ (replayFinal, _) <- runEff . runPrim . runTime . runTrace . runLLMReplay idx $ demoPipeline demoArticle+ replayCalls <- readIORef calls+ replayFinal @?= liveFinal+ replayCalls @?= 0+ ]++-- | A dependent two-stage pipeline: draft a summary of the article, then critique+-- the draft. The second stage's request depends on the first stage's output, so+-- replay must reproduce stage one exactly to reach the same stage-two request.+twoStage :: (LLM :> es) => Text -> Eff es (Text, Text)+twoStage article = do+ draft <- responseText <$> complete stubModel (ctxFor ("draft: " <> article)) optsFor+ crit <- responseText <$> complete stubModel (ctxFor ("critique: " <> draft)) optsFor+ pure (draft, crit)++-- | The deterministic stub: a draft request yields a fixed summary; a critique+-- request echoes the draft it was given. Deterministic per request, so replay is+-- well-defined.+responder :: Context -> Response+responder c =+ let t = lastUserText c+ in if "draft:" `T.isPrefixOf` t+ then mkResponse "a tight one-line summary"+ else mkResponse ("looks good: " <> T.drop (T.length "critique: ") t)++-- | The EP-6 golden fixture and its pinned digest, copied here so the two plans+-- are proven byte-for-byte identical.+fixModel :: Model+fixModel = _Model & #modelId .~ "claude-sonnet-4-6" & #provider .~ "anthropic"++fixCtx :: Context+fixCtx =+ _Context+ & #systemPrompt .~ Just "You are helpful."+ & #messages .~ V.singleton (user "ping")++fixOpts :: Options+fixOpts = _Options & #temperature .~ Just 0.0 & #maxTokens .~ Just 1024++pinnedKey :: Text+pinnedKey = "30b2015562ec8b5cd4fdb64c7cc671c84f56f80d24891deec6676c521f008113"++-- ---------------------------------------------------------------------------+-- Shared helpers+-- ---------------------------------------------------------------------------++-- | Run a two-stage traced computation: a program span containing two module+-- spans, each issuing one stubbed LM call (captured automatically by 'tracedLLM').+buildTree :: IO TraceTree+buildTree = do+ let resp = mkResponse "ok"+ ((), tree) <-+ runEff . runPrim . runTime . runTrace . runFixedLLM resp . tracedLLM $+ withSpan ProgramSpan "summarize-and-critique" $ do+ _ <- withSpan ModuleSpan "predict:Draft" (complete stubModel (ctxFor "draft this") optsFor)+ _ <- withSpan ModuleSpan "predict:Critique" (complete stubModel (ctxFor "critique that") optsFor)+ pure ()+ pure tree++spanAt :: TraceTree -> SpanId -> Span+spanAt tree sid = spans tree Map.! sid++isObject :: Value -> Bool+isObject v = case v of+ Object {} -> True+ _ -> False++-- ---------------------------------------------------------------------------+-- A small generator of random trees for the round-trip property+-- ---------------------------------------------------------------------------++baseTime :: UTCTime+baseTime = read "2026-06-08 00:00:00 UTC"++-- | A bare shape: a node kind and its children shapes.+data Shape = Shape SpanKind [Shape]++genShape :: Int -> Gen Shape+genShape d = do+ k <- elements [ProgramSpan, ModuleSpan, CombinatorSpan, LlmCallSpan]+ kids <-+ if d <= 0+ then pure []+ else do+ n <- choose (0, 3)+ vectorOf n (genShape (d - 1))+ pure (Shape k kids)++-- | Generate a random (single-rooted) tree of up to depth 3, with sequential span+-- ids, integer-second timestamps (exact JSON round-trip), and LM-call leaves+-- carrying a cacheKey + response so the serialized form exercises every field.+genTree :: Gen TraceTree+genTree = do+ shape <- genShape 3+ let (ss, _) = flattenShape Nothing 0 shape+ m = Map.fromList [(spanId s, s) | s <- ss]+ rootId = spanId (firstSpan ss)+ pure (TraceTree rootId m)+ where+ firstSpan (s : _) = s+ firstSpan [] = error "genTree: empty"++-- | Flatten a shape into spans, assigning ids/parents/times from a running+-- counter. Returns the produced spans (root first) and the next free counter.+flattenShape :: Maybe SpanId -> Int -> Shape -> ([Span], Int)+flattenShape par n (Shape k kids) =+ let sid = SpanId ("span-" <> T.pack (show n))+ s =+ Span+ { spanId = sid,+ parent = par,+ kind = k,+ label = "lbl-" <> T.pack (show n),+ startedAt = addUTCTime (fromIntegral (2 * n)) baseTime,+ endedAt = Just (addUTCTime (fromIntegral (2 * n + 1)) baseTime),+ attrs = attrsFor k n+ }+ (kidSpans, n') = foldl step ([], n + 1) kids+ step (acc, cnt) kidShape =+ let (ks, cnt') = flattenShape (Just sid) cnt kidShape+ in (acc ++ ks, cnt')+ in (s : kidSpans, n')++-- | LM-call leaves carry a key + response (exercising replayIndex and the Value+-- fields); structural nodes are empty.+attrsFor :: SpanKind -> Int -> SpanAttrs+attrsFor LlmCallSpan n =+ emptyAttrs+ { model = Just "m",+ cacheKey = Just ("key-" <> T.pack (show n)),+ response = Just (object ["text" .= ("resp-" <> T.pack (show n))]),+ inputTokens = Just (fromIntegral n),+ costUsd = Just (read "0.001")+ }+attrsFor _ _ = emptyAttrs++-- ---------------------------------------------------------------------------+-- EP-16: node identity, correlation, and the feedback channel+-- ---------------------------------------------------------------------------++-- | A self-looping record so @predict cellSig :: Program Cell Cell@ composes with+-- itself to make a multi-node program.+newtype Cell = Cell {cell :: Text}+ deriving stock (Generic, Show, Eq)++instance ToSchema Cell++instance FromModel Cell++instance ToPrompt Cell++instance Validatable Cell++cellSig :: Signature Cell Cell+cellSig = mkSignature "Echo the cell"++-- | A two-@Predict@ chain: @Compose (Predict) (Predict)@.+chain2 :: Program Cell Cell+chain2 = predict cellSig >>> predict cellSig++-- | A fixed response in fallback @[[ ## field ## ]]@ form, decoded by the lenient+-- parser into @Cell "echoed"@.+cellResp :: Response+cellResp = mkResponse "[[ ## cell ## ]]\nechoed\n[[ ## completed ## ]]\n"++chkLen :: Program a b -> IO ()+chkLen p = length (programNodePaths p) @?= length (foldParams p)++-- | A total ensemble reducer (the program is only shape-inspected, never run).+firstOr :: [Cell] -> Cell+firstOr (c : _) = c+firstOr [] = Cell ""++-- | The @LlmCallSpan@s of a tree, in start order.+llmSpansInOrder :: TraceTree -> [Span]+llmSpansInOrder t = sortOn startedAt [s | s <- Map.elems (spans t), kind s == LlmCallSpan]++-- | Run a program under 'runProgramTraced' with the node-aware stack, returning the+-- decoded result and the captured tree (or a typed error).+runTraced :: Program Cell Cell -> Cell -> IO (Either ShikumiError (Cell, TraceTree))+runTraced prog input =+ runEff+ . runErrorNoCallStack @ShikumiError+ . runPrim+ . runTime+ . runTrace+ . runCurrentNode+ . runFixedLLM cellResp+ . tracedNodeLLM+ $ runProgramTraced prog input++-- | M1: 'NodePath' identity, its agreement with @foldParams@, field metadata, and+-- JSON round-trips.+nodeTests :: TestTree+nodeTests =+ testGroup+ "node"+ [ testCase "programNodePaths length agrees with foldParams across shapes" $ do+ chkLen chain2+ chkLen (majorityVote 3 (TempFixed []) (predict cellSig))+ chkLen (majorityVoteBy 2 (TempFixed []) firstOr (predict cellSig))+ chkLen (parallel2 (predict cellSig) (predict cellSig)),+ testCase "programNodePaths yields the expected structural paths for a chain" $+ programNodePaths chain2 @?= [NodePath [StepComposeL], NodePath [StepComposeR]],+ testCase "nodeFields returns each node's input/output field names" $ do+ let nf = map snd (nodeFields chain2)+ map inputFieldNames nf @?= [["cell"], ["cell"]]+ map outputFieldNames nf @?= [["cell"], ["cell"]],+ testCase "mapParamsAt k edits the node at programNodePaths !! k (index law)" $ do+ let edited = mapParamsAt 1 (\ps -> ps {instructionOverride = Just "x"}) chain2+ map instructionOverride (foldParams edited) @?= [Nothing, Just "x"],+ testCase "NodePath and a nodePath-bearing SpanAttrs round-trip as JSON" $ do+ let np = NodePath [StepComposeL, StepEnsemble 2]+ decode (encode np) @?= Just np+ let a = emptyAttrs {nodePath = Just np}+ decode (encode a) @?= Just a+ ]++-- | M2: each model-call span is tagged with its issuing node's 'NodePath', and the+-- nodePath-bearing tree round-trips through the store at formatVersion 2.+correlateTests :: TestTree+correlateTests =+ testGroup+ "correlate"+ [ testCase "each llm-call span carries its node's path, in programNodePaths order" $ do+ res <- runTraced chain2 (Cell "x")+ case res of+ Left e -> assertFailure ("traced run failed: " <> show e)+ Right (out, tree) -> do+ out @?= Cell "echoed"+ let llms = llmSpansInOrder tree+ length llms @?= 2+ let paths = programNodePaths chain2+ map (nodePath . attrs) llms @?= [Just (paths !! 0), Just (paths !! 1)],+ testCase "a nodePath-bearing tree round-trips through the store (formatVersion 2)" $+ withSystemTempDirectory "shikumi-trace-node" $ \dir -> do+ res <- runTraced chain2 (Cell "x")+ case res of+ Left e -> assertFailure ("traced run failed: " <> show e)+ Right (_, tree) -> do+ let p = dir <> "/trace.json"+ writeTraceFile p tree+ loaded <- readTraceFile p+ loaded @?= Right tree+ ]++-- | M3: the per-node feedback channel round-trips writes and reads.+feedbackTests :: TestTree+feedbackTests =+ testGroup+ "feedback"+ [ testCase "feedbackFor returns critiques in attach order; untargeted node is empty" $ do+ let p1 = NodePath [StepComposeL]+ p2 = NodePath [StepComposeR]+ p3 = NodePath [StepFMap]+ (_, fb) <-+ runEff . runPrim . runFeedback $ do+ attachFeedback p1 "first"+ attachFeedback p1 "second"+ attachFeedback p2 "only"+ feedbackFor p1 fb @?= ["first", "second"]+ feedbackFor p2 fb @?= ["only"]+ feedbackFor p3 fb @?= [],+ testCase "FeedbackLog round-trips as JSON" $ do+ (_, fb) <-+ runEff . runPrim . runFeedback $+ attachFeedback (NodePath [StepComposeL]) "c"+ decode (encode (fb :: FeedbackLog)) @?= Just fb+ ]
+ test/TraceFixtures.hs view
@@ -0,0 +1,103 @@+-- | Shared, network-free fixtures for the EP-7 trace tests: a stub 'Model', a+-- request builder, canned responses, and two base @LLM@ interpreters — a /fixed/+-- one that returns the same response for every call, and a /keyed/ one that picks+-- a response from the request 'Context' (so a two-stage pipeline gets distinct+-- responses, hence distinct cache keys).+module TraceFixtures+ ( stubModel,+ ctxFor,+ optsFor,+ mkResponse,+ responseText,+ lastUserText,+ runFixedLLM,+ runKeyedLLM,+ runKeyedCountingLLM,+ )+where++import Baikai+ ( Api (Custom),+ AssistantContent (..),+ Context,+ Message (UserMessage),+ Model,+ Options,+ Response,+ TextContent (..),+ UserContent (UserText),+ user,+ _Context,+ _Model,+ _Options,+ _Response,+ _TextContent,+ )+import Control.Lens ((&), (.~), (^.))+import Data.Generics.Labels ()+import Data.IORef (IORef, atomicModifyIORef')+import Data.Text (Text)+import Data.Vector qualified as V+import Effectful (Eff, IOE, liftIO, (:>))+import Effectful.Dispatch.Dynamic (interpret)+import Shikumi.LLM (LLM (..))++-- | A stub model with real-looking routing identity, so LM-call span labels read+-- @stub/stub-model@.+stubModel :: Model+stubModel =+ _Model+ & #provider .~ "stub"+ & #modelId .~ "stub-model"+ & #api .~ Custom "stub"++-- | A one-user-turn request context carrying @t@ (varies the cache key per stage).+ctxFor :: Text -> Context+ctxFor t = _Context & #messages .~ V.singleton (user t)++-- | Default options.+optsFor :: Options+optsFor = _Options++-- | A response carrying @t@ as its single assistant text block, with small fixed+-- token counts and latency so span attributes are populated.+mkResponse :: Text -> Response+mkResponse t =+ _Response+ & #message . #content .~ V.singleton (AssistantText (_TextContent & #text .~ t))+ & #message . #usage . #inputTokens .~ 10+ & #message . #usage . #outputTokens .~ 5+ & #latencyMs .~ 12++-- | Concatenate the assistant text blocks of a response.+responseText :: Response -> Text+responseText resp =+ mconcat [t | AssistantText (TextContent t) <- V.toList (resp ^. #message . #content)]++-- | Concatenate the text of every user message in a request context.+lastUserText :: Context -> Text+lastUserText c =+ mconcat+ [ t+ | UserMessage up <- V.toList (c ^. #messages),+ UserText (TextContent t) <- V.toList (up ^. #content)+ ]++-- | A base @LLM@ interpreter that returns the same response for every completion.+runFixedLLM :: Response -> Eff (LLM : es) a -> Eff es a+runFixedLLM resp = interpret $ \_ -> \case+ Complete {} -> pure resp+ Stream {} -> pure []++-- | A base @LLM@ interpreter that picks a response from the request context.+runKeyedLLM :: (Context -> Response) -> Eff (LLM : es) a -> Eff es a+runKeyedLLM f = interpret $ \_ -> \case+ Complete _ c _ -> pure (f c)+ Stream {} -> pure []++-- | Like 'runKeyedLLM' but bumps a counter on every completion (used to assert how+-- many provider calls a live run made — and that a replay run makes zero).+runKeyedCountingLLM :: (IOE :> es) => IORef Int -> (Context -> Response) -> Eff (LLM : es) a -> Eff es a+runKeyedCountingLLM ref f = interpret $ \_ -> \case+ Complete _ c _ -> liftIO (atomicModifyIORef' ref (\n -> (n + 1, ()))) >> pure (f c)+ Stream {} -> pure []