diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/shikumi-trace.cabal b/shikumi-trace.cabal --- a/shikumi-trace.cabal +++ b/shikumi-trace.cabal @@ -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 diff --git a/src/Shikumi/Trace.hs b/src/Shikumi/Trace.hs --- a/src/Shikumi/Trace.hs +++ b/src/Shikumi/Trace.hs @@ -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) diff --git a/src/Shikumi/Trace/Demo.hs b/src/Shikumi/Trace/Demo.hs --- a/src/Shikumi/Trace/Demo.hs +++ b/src/Shikumi/Trace/Demo.hs @@ -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 diff --git a/src/Shikumi/Trace/Node.hs b/src/Shikumi/Trace/Node.hs --- a/src/Shikumi/Trace/Node.hs +++ b/src/Shikumi/Trace/Node.hs @@ -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 _) = [] diff --git a/src/Shikumi/Trace/Program.hs b/src/Shikumi/Trace/Program.hs --- a/src/Shikumi/Trace/Program.hs +++ b/src/Shikumi/Trace/Program.hs @@ -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 diff --git a/src/Shikumi/Trace/Store.hs b/src/Shikumi/Trace/Store.hs --- a/src/Shikumi/Trace/Store.hs +++ b/src/Shikumi/Trace/Store.hs @@ -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] diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -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. diff --git a/test/TraceFixtures.hs b/test/TraceFixtures.hs --- a/test/TraceFixtures.hs +++ b/test/TraceFixtures.hs @@ -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 []