packages feed

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 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 []