shikumi-trace 0.1.1.0 → 0.2.0.0
raw patch · 9 files changed
+289/−53 lines, 9 filesdep ~baikaidep ~shikumidep ~shikumi-cachePVP ok
version bump matches the API change (PVP)
Dependency ranges changed: baikai, shikumi, shikumi-cache, shikumi-trace
API changes (from Hackage documentation)
+ Shikumi.Trace.Store: minSupportedFormatVersion :: Int
- Shikumi.Trace.Store: replayIndex :: TraceTree -> Map CacheKey Value
+ Shikumi.Trace.Store: replayIndex :: TraceTree -> Either Text (Map CacheKey Value)
Files
- CHANGELOG.md +18/−0
- shikumi-trace.cabal +9/−9
- src/Shikumi/Trace.hs +44/−12
- src/Shikumi/Trace/Demo.hs +6/−4
- src/Shikumi/Trace/Node.hs +1/−1
- src/Shikumi/Trace/Program.hs +26/−7
- src/Shikumi/Trace/Store.hs +38/−12
- test/Main.hs +134/−8
- test/TraceFixtures.hs +13/−0
CHANGELOG.md view
@@ -2,6 +2,24 @@ ## Unreleased +## 0.2.0.0 - 2026-07-05++### Added++- `minSupportedFormatVersion` documents the oldest trace file format accepted by+ this build.++### Changed++- **BREAKING** `replayIndex` now returns `Either Text (Map CacheKey Value)` and+ fails closed when a trace records conflicting responses for the same cache key.+- Trace span mutation now fails loudly on unsupported concurrent stack use instead+ of silently corrupting the tree; `runTrace`/`tracedLLM` remain sequential.+- Retry spans now count traced retry attempts, trace rendering handles multiple+ roots, and child ordering is stable by numeric `span-N` ids.+- Refreshed internal `shikumi` and `shikumi-cache` bounds for the current package+ set.+ ## 0.1.1.0 - 2026-06-28 ### Changed
shikumi-trace.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.4 name: shikumi-trace-version: 0.1.1.0+version: 0.2.0.0 synopsis: Hierarchical tracing, observability, and deterministic replay for shikumi (EP-7) @@ -51,7 +51,7 @@ build-depends: , aeson- , baikai >=0.2 && <0.3+ , baikai >=0.3 && <0.4 , base >=4.20 && <5 , bytestring , containers@@ -60,8 +60,8 @@ , generic-lens , lens ^>=5.3 , scientific- , shikumi ^>=0.2.0.0- , shikumi-cache ^>=0.1.1.0+ , shikumi ^>=0.3.0.0+ , shikumi-cache ^>=0.1.2.0 , text ^>=2.1 , time , vector@@ -73,7 +73,7 @@ ghc-options: -threaded -with-rtsopts=-N build-depends: , base- , shikumi-trace ^>=0.1.1.0+ , shikumi-trace ^>=0.2.0.0 test-suite shikumi-trace-test import: common-options@@ -84,7 +84,7 @@ other-modules: TraceFixtures build-depends: , aeson- , baikai >=0.2 && <0.3+ , baikai >=0.3 && <0.4 , base , bytestring , containers@@ -92,9 +92,9 @@ , generic-lens , lens ^>=5.3 , QuickCheck- , shikumi ^>=0.2.0.0- , shikumi-cache ^>=0.1.1.0- , shikumi-trace ^>=0.1.1.0+ , shikumi ^>=0.3.0.0+ , shikumi-cache ^>=0.1.2.0+ , shikumi-trace ^>=0.2.0.0 , tasty , tasty-hunit , tasty-quickcheck
src/Shikumi/Trace.hs view
@@ -79,7 +79,6 @@ import Effectful.Prim.IORef ( IORef, atomicModifyIORef',- modifyIORef', newIORef, readIORef, )@@ -91,6 +90,7 @@ import Shikumi.LLM (LLM (..), complete, stream) import Shikumi.Trace.Node (NodePath (..)) import Shikumi.Trace.ResponseJSON ()+import Text.Read (readMaybe) -- --------------------------------------------------------------------------- -- Span and tree types@@ -184,13 +184,16 @@ deriving anyclass (ToJSON, FromJSON) -- | The children of a span, in creation order (sorted by start time, ties broken--- by span id).+-- by numeric @span-N@ id when possible). childrenOf :: TraceTree -> SpanId -> [SpanId] childrenOf t sid = map (^. #spanId) $- sortOn (\s -> (s ^. #startedAt, s ^. #spanId)) $+ sortOn (\s -> (s ^. #startedAt, spanOrdKey (s ^. #spanId))) $ [s | s <- Map.elems (t ^. #spans), (s ^. #parent) == Just sid] +spanOrdKey :: SpanId -> (Maybe Int, Text)+spanOrdKey (SpanId t) = (T.stripPrefix "span-" t >>= readMaybe . T.unpack, t)+ -- --------------------------------------------------------------------------- -- The effect -- ---------------------------------------------------------------------------@@ -246,6 +249,11 @@ -- 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'.+--+-- The span stack is sequential: use 'Shikumi.Program.runProgram' with+-- 'tracedLLM', not 'Shikumi.Program.runProgramConc'. State writes are atomic so a+-- concurrent send cannot tear the span map, but a concurrent close that violates+-- stack order now fails loudly instead of recording a silently wrong tree. runTrace :: (Prim :> es, Time :> es) => Eff (Trace : es) a -> Eff es (a, TraceTree) runTrace act = do st <- newTraceState@@ -280,26 +288,42 @@ par <- safeHead <$> readIORef (st ^. #stack) now <- getCurrentTime let s = Span sid par k lbl now Nothing emptyAttrs- modifyIORef' (st ^. #spans) (at sid ?~ s)- modifyIORef' (st ^. #stack) (sid :)+ atomicModifyIORef' (st ^. #spans) (\m -> (m & at sid ?~ s, ()))+ atomicModifyIORef' (st ^. #stack) (\stk -> (sid : stk, ())) case par of- Nothing -> modifyIORef' (st ^. #root) (Just . fromMaybe sid)+ Nothing -> atomicModifyIORef' (st ^. #root) (\r -> (Just (fromMaybe sid r), ())) Just _ -> pure () pure sid --- | Close a span: stamp its 'endedAt' and pop it off the stack.+-- | Close a span: stamp its 'endedAt' and pop it off the stack. The pop is+-- verified because 'withSpan' brackets guarantee LIFO close order in sequential+-- execution. Popping anything else means unsupported concurrent trace mutation,+-- so fail loudly instead of corrupting the trace. closeSpan :: (Prim :> es, Time :> es) => TraceState -> SpanId -> Eff es () closeSpan st sid = do now <- getCurrentTime- modifyIORef' (st ^. #spans) (ix sid . #endedAt ?~ now)- modifyIORef' (st ^. #stack) (drop 1)+ atomicModifyIORef' (st ^. #spans) (\m -> (m & ix sid . #endedAt ?~ now, ()))+ popped <- atomicModifyIORef' (st ^. #stack) $ \case+ top : rest -> (rest, Just top)+ [] -> ([], Nothing)+ case popped of+ Just top | top == sid -> pure ()+ _ ->+ error+ ( "Shikumi.Trace.runTrace: span stack corrupted (closing "+ <> show sid+ <> " but popped "+ <> show popped+ <> "). runTrace supports sequential execution only; do not compose "+ <> "tracedLLM/tracedNodeLLM with runProgramConc."+ ) -- | 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 (st ^. #stack) case stk of- (sid : _) -> modifyIORef' (st ^. #spans) (ix sid . #attrs %~ f)+ (sid : _) -> atomicModifyIORef' (st ^. #spans) (\m -> (m & ix sid . #attrs %~ f, ())) [] -> pure () -- | Freeze the building state into an immutable 'TraceTree'.@@ -322,6 +346,10 @@ -- 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').+--+-- This capture uses 'runTrace'\'s sequential span stack. Compose it with+-- 'Shikumi.Program.runProgram'; concurrent program execution is intentionally+-- outside the current trace builder's contract. tracedLLM :: (Trace :> es, LLM :> es) => Eff es a -> Eff es a tracedLLM = interpose $ \_ -> \case Complete m c o -> withSpan LlmCallSpan (llmLabel m) $ do@@ -342,7 +370,7 @@ & #provider ?~ (m ^. #provider) & #prompt ?~ requestToCanonicalValue m c o & #response ?~ toJSON resp- & #latencyMs ?~ (resp ^. #latencyMs)+ & #latencyMs ?~ fromIntegral (resp ^. #latencyMs) & #inputTokens ?~ (resp ^. #message . #usage . #inputTokens) & #outputTokens ?~ (resp ^. #message . #usage . #outputTokens) & #costUsd ?~ realToFrac (resp ^. #message . #usage . #cost . #usd :: Rational)@@ -365,8 +393,12 @@ renderTree :: TraceTree -> Text renderTree t | Map.null (t ^. #spans) = "(empty trace)\n"- | otherwise = T.concat (go (0 :: Int) (t ^. #root))+ | otherwise = T.concat (concatMap (go (0 :: Int)) roots) where+ roots =+ map (^. #spanId) $+ sortOn (\s -> (s ^. #startedAt, spanOrdKey (s ^. #spanId))) $+ [s | s <- Map.elems (t ^. #spans), (s ^. #parent) == Nothing] go depth sid = case Map.lookup sid (t ^. #spans) of Nothing -> [] Just s -> line depth s : concatMap (go (depth + 1)) (childrenOf t sid)
src/Shikumi/Trace/Demo.hs view
@@ -133,10 +133,12 @@ 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"+ case replayIndex tree of+ Left err -> TIO.putStrLn ("replay index error: " <> err)+ Right idx -> do+ (critique, _) <- runEff . runPrim . runTime . runTrace . runLLMReplay idx $ demoPipeline demoArticle+ TIO.putStrLn ("FINAL: " <> critique)+ TIO.putStrLn "provider calls: 0" -- --------------------------------------------------------------------------- -- helpers
src/Shikumi/Trace/Node.hs view
@@ -94,7 +94,7 @@ 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 (MajorityVote _ _ _ p) = go (StepMajorityVote : prefix) p go prefix (Ensemble ps _) = concat (zipWith (\i p -> go (StepEnsemble i : prefix) p) [0 ..] ps) go _ (Embed _) = []
src/Shikumi/Trace/Program.hs view
@@ -46,7 +46,7 @@ import Data.Generics.Labels () import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, (:>)) import Effectful.Dispatch.Dynamic (interpose, interpret, localSeqUnlift, send)-import Effectful.Error.Static (Error)+import Effectful.Error.Static (Error, catchError, throwError) import Effectful.Exception (bracket_) import Effectful.Prim (Prim) import Effectful.Prim.IORef (newIORef, readIORef, writeIORef)@@ -67,8 +67,6 @@ Validate ), acceptOrReject,- modal,- retryWith, runProgram, sampleTemps, withSampleTemp,@@ -77,6 +75,7 @@ ( SpanKind (CombinatorSpan, LlmCallSpan, ModuleSpan), Trace, annotateSpan,+ bumpRetry, llmAttrs, llmLabel, withSpan,@@ -110,6 +109,10 @@ -- 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.+--+-- The cell is dynamically scoped for sequential execution. Do not compose+-- 'runCurrentNode'\/'tracedNodeLLM' with 'Shikumi.Program.runProgramConc': sibling+-- branches would share one mutable current-node slot and could mis-scope paths. runCurrentNode :: (Prim :> es) => Eff (CurrentNode : es) a -> Eff es a runCurrentNode act = do ref <- newIORef Nothing@@ -172,16 +175,32 @@ 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)+ withSpan CombinatorSpan "Retry" (tracedRetry (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)+ withSpan CombinatorSpan "RetryWhen" (tracedRetry (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 =+ go prefix (MajorityVote k sched reduce p) i = withSpan CombinatorSpan "MajorityVote" $- modal <$> traverse (\mt -> withSampleTemp mt (go (StepMajorityVote : prefix) p i)) (sampleTemps (max 1 k) sched)+ reduce <$> traverse (\mt -> withSampleTemp mt (go (StepMajorityVote : prefix) p i)) (sampleTemps 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)++tracedRetry ::+ (Trace :> es, Error ShikumiError :> es) =>+ (Program x y -> x -> Eff es y) ->+ (ShikumiError -> Bool) ->+ Int ->+ Program x y ->+ x ->+ Eff es y+tracedRetry run ok n p i = attempt (max 1 n)+ where+ attempt left =+ run p i `catchError` \_cs e ->+ if ok e && left > 1+ then bumpRetry >> attempt (left - 1)+ else throwError e
src/Shikumi/Trace/Store.hs view
@@ -16,6 +16,7 @@ -- 'Shikumi.Cache.Key.CacheKey' mapped to its recorded response JSON. module Shikumi.Trace.Store ( TraceFile (..),+ minSupportedFormatVersion, currentFormatVersion, writeTraceFile, readTraceFile,@@ -27,13 +28,14 @@ import Data.Aeson (FromJSON, ToJSON, Value, eitherDecode, encode) import Data.ByteString.Lazy qualified as BL import Data.Generics.Labels ()+import Data.List.NonEmpty qualified as NE 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 Shikumi.Trace (SpanId (..), SpanKind (LlmCallSpan), TraceTree (..)) import System.Directory (renameFile) -- | The on-disk wrapper: a format version plus the tree. Bump@@ -51,6 +53,11 @@ currentFormatVersion :: Int currentFormatVersion = 2 +-- | The oldest trace format this build still reads. v1→v2 was additive (the+-- optional @SpanAttrs.nodePath@ field), so v1 files decode without migration.+minSupportedFormatVersion :: Int+minSupportedFormatVersion = 1+ -- | 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 ()@@ -67,11 +74,13 @@ pure $ case eitherDecode bs of Left e -> Left ("trace parse error: " <> T.pack e) Right tf- | formatVersion tf /= currentFormatVersion ->+ | formatVersion tf < minSupportedFormatVersion || formatVersion tf > currentFormatVersion -> Left ( "unsupported trace formatVersion " <> T.pack (show (formatVersion tf))- <> " (expected "+ <> " (supported: "+ <> T.pack (show minSupportedFormatVersion)+ <> ".." <> T.pack (show currentFormatVersion) <> ")" )@@ -79,13 +88,30 @@ -- | 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+-- contributes nothing. Duplicate keys are legal only when every occurrence+-- recorded the same response; differing responses mean replay is not+-- deterministic, so index construction fails closed.+replayIndex :: TraceTree -> Either Text (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]- ]+ case conflicts of+ [] -> Right (Map.map (snd . NE.head) grouped)+ cs -> Left (T.intercalate "; " (map describe cs))+ where+ occurrences =+ [ (CacheKey ck, (s ^. #spanId, v))+ | s <- Map.elems (spans t),+ (s ^. #kind) == LlmCallSpan,+ Just ck <- [s ^. #attrs . #cacheKey],+ Just v <- [s ^. #attrs . #response]+ ]+ grouped = Map.fromListWith (<>) [(k, NE.singleton sv) | (k, sv) <- occurrences]+ conflicts =+ [ (k, NE.toList svs)+ | (k, svs) <- Map.toList grouped,+ length (NE.nub (fmap snd svs)) > 1+ ]+ describe (CacheKey k, svs) =+ "replay index conflict: cache key "+ <> k+ <> " was recorded with differing responses by spans "+ <> T.intercalate ", " [sid | (SpanId sid, _) <- svs]
test/Main.hs view
@@ -22,7 +22,7 @@ import Shikumi.Adapter (ToPrompt) import Shikumi.Cache.Key (CacheKey (..)) import Shikumi.Cache.Key qualified as Key-import Shikumi.Combinator (majorityVote, majorityVoteBy, parallel2, (>>>))+import Shikumi.Combinator (majorityVote, majorityVoteBy, parallel2, retry, (>>>)) import Shikumi.Effect.Time (runTime) import Shikumi.Error (ShikumiError) import Shikumi.LLM (LLM, complete)@@ -34,6 +34,7 @@ TempSchedule (TempFixed), foldParams, mapParamsAt,+ runProgram, ) import Shikumi.Schema (FromModel, ToSchema, Validatable) import Shikumi.Signature (Signature, mkSignature)@@ -73,6 +74,7 @@ import Shikumi.Trace.Store ( TraceFile (..), currentFormatVersion,+ minSupportedFormatVersion, readTraceFile, replayIndex, writeTraceFile,@@ -89,6 +91,7 @@ responseText, runFixedLLM, runKeyedCountingLLM,+ runSequencedLLM, stubModel, ) @@ -142,7 +145,18 @@ 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)+ assertBool "module lines present" (T.count "module" out == 2),+ testCase "renderTree shows every top-level root" $ do+ ((), tree) <-+ runEff . runPrim . runTime . runTrace $ do+ withSpan ProgramSpan "first-root" (pure ())+ withSpan ProgramSpan "second-root" (pure ())+ let out = renderTree tree+ assertBool "first root is rendered" ("first-root" `T.isInfixOf` out)+ assertBool "second root is rendered" ("second-root" `T.isInfixOf` out),+ testCase "childrenOf orders 10+ siblings numerically" $ do+ let tree = numericSiblingTree+ childrenOf tree (SpanId "span-1") @?= map (SpanId . ("span-" <>) . T.pack . show) ([2 .. 12] :: [Int]) ] -- ---------------------------------------------------------------------------@@ -175,12 +189,45 @@ 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 "reading a formatVersion 1 file succeeds" $+ withSystemTempDirectory "shikumi-trace" $ \dir -> do+ tree <- buildTree+ let p = dir <> "/v1.json"+ minSupportedFormatVersion @?= 1+ BL.writeFile p (encode (TraceFile 1 tree))+ res <- readTraceFile p+ res @?= Right tree, testCase "replayIndex maps each llm-call cacheKey to its response" $ do tree <- buildTree- let idx = replayIndex tree+ idx <- replayIndexOrFail 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 "duplicate keys with equal responses dedupe" $ do+ let tree = duplicateKeyTree (object ["text" .= ("same" :: Text)]) (object ["text" .= ("same" :: Text)])+ idx <- replayIndexOrFail tree+ Map.size idx @?= 1,+ testCase "duplicate keys with conflicting responses fail closed" $ do+ let tree = duplicateKeyTree (object ["text" .= ("first" :: Text)]) (object ["text" .= ("second" :: Text)])+ case replayIndex tree of+ Left msg -> do+ assertBool "message names the key" ("dup-key" `T.isInfixOf` msg)+ assertBool "message names span-1" ("span-1" `T.isInfixOf` msg)+ assertBool "message names span-2" ("span-2" `T.isInfixOf` msg)+ Right _ -> assertFailure "expected conflicting duplicate keys to fail closed",+ testCase "majorityVote over identical requests replays" $ do+ (live, tree) <-+ runEff . runPrim . runTime . runTrace . runFixedLLM cellResp . tracedLLM . runErrorNoCallStack @ShikumiError $+ runProgram (majorityVote 3 (TempFixed []) (predict cellSig)) (Cell "x")+ case live of+ Left e -> assertFailure ("live majorityVote failed: " <> show e)+ Right liveOut -> do+ idx <- replayIndexOrFail tree+ Map.size idx @?= 1+ replayed <-+ runEff . runLLMReplay idx . runErrorNoCallStack @ShikumiError $+ runProgram (majorityVote 3 (TempFixed []) (predict cellSig)) (Cell "x")+ replayed @?= Right liveOut, testCase "cacheKey reproduces EP-6's pinned digest (integration point #7)" $ do let CacheKey hex = Key.cacheKey fixModel fixCtx fixOpts hex @?= pinnedKey@@ -210,7 +257,7 @@ let p = dir <> "/trace.json" writeTraceFile p tree Right tree' <- readTraceFile p- let idx = replayIndex tree'+ idx <- replayIndexOrFail tree' -- (4) replay the SAME pipeline; the counting provider must not be hit. writeIORef calls 0 replayed <- runEff . runLLMReplay idx $ twoStage "the article"@@ -227,7 +274,7 @@ let p = dir <> "/trace.json" writeTraceFile p tree Right tree' <- readTraceFile p- let idx = replayIndex tree'+ idx <- replayIndexOrFail tree' -- A different article => a first-stage key that was never recorded. res <- try (runEff . runLLMReplay idx $ twoStage "a different article") case res of@@ -260,7 +307,7 @@ let p = dir <> "/trace.json" writeTraceFile p tree Right tree' <- readTraceFile p- let idx = replayIndex tree'+ idx <- replayIndexOrFail tree' writeIORef calls 0 (replayFinal, _) <- runEff . runPrim . runTime . runTrace . runLLMReplay idx $ demoPipeline demoArticle replayCalls <- readIORef calls@@ -302,7 +349,7 @@ fixOpts = _Options & #temperature .~ Just 0.0 & #maxTokens .~ Just 1024 pinnedKey :: Text-pinnedKey = "30b2015562ec8b5cd4fdb64c7cc671c84f56f80d24891deec6676c521f008113"+pinnedKey = "b31fd70140abbd0198c6b7caec748a8389bf93be909164bdcc340731b7032564" -- --------------------------------------------------------------------------- -- Shared helpers@@ -329,6 +376,65 @@ Object {} -> True _ -> False +replayIndexOrFail :: TraceTree -> IO (Map.Map CacheKey Value)+replayIndexOrFail tree =+ either (assertFailure . T.unpack) pure (replayIndex tree)++duplicateKeyTree :: Value -> Value -> TraceTree+duplicateKeyTree firstResp secondResp =+ let rootSpan =+ Span+ { spanId = SpanId "span-0",+ parent = Nothing,+ kind = ProgramSpan,+ label = "root",+ startedAt = baseTime,+ endedAt = Just (addUTCTime 1 baseTime),+ attrs = emptyAttrs+ }+ child :: Int -> Value -> Span+ child n resp =+ Span+ { spanId = SpanId ("span-" <> T.pack (show n)),+ parent = Just (SpanId "span-0"),+ kind = LlmCallSpan,+ label = "llm-call",+ startedAt = addUTCTime (fromIntegral n) baseTime,+ endedAt = Just (addUTCTime (fromIntegral n + 1) baseTime),+ attrs =+ emptyAttrs+ { cacheKey = Just "dup-key",+ response = Just resp+ }+ }+ ss = [rootSpan, child 1 firstResp, child 2 secondResp]+ in TraceTree (SpanId "span-0") (Map.fromList [(spanId s, s) | s <- ss])++numericSiblingTree :: TraceTree+numericSiblingTree =+ let rootSpan =+ Span+ { spanId = SpanId "span-1",+ parent = Nothing,+ kind = ProgramSpan,+ label = "root",+ startedAt = baseTime,+ endedAt = Just (addUTCTime 1 baseTime),+ attrs = emptyAttrs+ }+ child n =+ Span+ { spanId = SpanId ("span-" <> T.pack (show n)),+ parent = Just (SpanId "span-1"),+ kind = ModuleSpan,+ label = "child-" <> T.pack (show n),+ startedAt = baseTime,+ endedAt = Just (addUTCTime 1 baseTime),+ attrs = emptyAttrs+ }+ ss = rootSpan : map child ([2 .. 12] :: [Int])+ in TraceTree (SpanId "span-1") (Map.fromList [(spanId s, s) | s <- ss])+ -- --------------------------------------------------------------------------- -- A small generator of random trees for the round-trip property -- ---------------------------------------------------------------------------@@ -505,7 +611,27 @@ let p = dir <> "/trace.json" writeTraceFile p tree loaded <- readTraceFile p- loaded @?= Right tree+ loaded @?= Right tree,+ testCase "Retry re-attempts are counted on the Retry span" $ do+ responses <- newIORef [mkResponse "unparseable", cellResp]+ res <-+ runEff+ . runErrorNoCallStack @ShikumiError+ . runPrim+ . runTime+ . runTrace+ . runCurrentNode+ . runSequencedLLM responses+ . tracedNodeLLM+ $ runProgramTraced (retry 2 (predict cellSig)) (Cell "x")+ case res of+ Left e -> assertFailure ("retry traced run failed: " <> show e)+ Right (out, tree) -> do+ out @?= Cell "echoed"+ let retrySpans = [s | s <- Map.elems (spans tree), label s == "Retry"]+ case retrySpans of+ [s] -> retries (attrs s) @?= 1+ _ -> assertFailure ("expected one Retry span, found " <> show (length retrySpans)) ] -- | M3: the per-node feedback channel round-trips writes and reads.
test/TraceFixtures.hs view
@@ -13,6 +13,7 @@ runFixedLLM, runKeyedLLM, runKeyedCountingLLM,+ runSequencedLLM, ) where @@ -100,4 +101,16 @@ 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 []++-- | A base @LLM@ interpreter that returns and removes the next response from an+-- 'IORef' list on each completion. Used for retry tests that need attempt one+-- and attempt two to receive different payloads.+runSequencedLLM :: (IOE :> es) => IORef [Response] -> Eff (LLM : es) a -> Eff es a+runSequencedLLM ref = interpret $ \_ -> \case+ Complete {} ->+ liftIO $+ atomicModifyIORef' ref $ \case+ r : rs -> (rs, r)+ [] -> ([], mkResponse "runSequencedLLM: exhausted") Stream {} -> pure []