diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,35 @@ ## [Unreleased] +## [baikai 0.10.1.0] - 2026-11-03 + +### Added + +- `Baikai.Trace.withTraceFold` and `withTraceFoldWith` own tracing across the + entire consumer fold; `withTraceList` and `withTraceListWith` collect events. + A healthy sink receives abort evidence before cancellation propagates, + including cancellation between stream steps. + +### Fixed + +- `withTrace`, `withTraceWith`, and the traced call in `runRequestWith` and + `runRequestWithRegistry` now finalize synchronously on cancellation. Cleanup + resumes an interrupted terminal drain within the same one-second budget, + preserves the original exception, and never duplicates terminal evidence. + General-stream abandonment keeps its garbage-collection-eventual contract. + +## [baikai-effectful 0.5.0.0] - 2026-11-03 + +### Added + +- `streamCollectTraced` delegates traced event collection to core with bounded + cleanup on cancellation. Re-export `TraceSink` for operation signatures. + +### Changed (breaking) + +- `Baikai` gains `StreamCollectTraced`. Third-party interpreters must handle + this constructor. The binding requires `baikai ^>=0.10.1`. + ## [baikai 0.10.0.0] - 2026-10-25 ### Added @@ -108,6 +137,8 @@ - Layered job configuration accepts `thinking "none"` through the shared canonical parser, keeping explicit off distinct from an absent preference. +- The UTF-8 stdin test runs in a child process under an ASCII locale, avoiding + shared-stdin descriptor races with parallel subprocess tests in the tarball. ### Changed diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@ data Baikai :: Effect where Complete :: Model -> Context -> Options -> Baikai m Response StreamCollect :: Model -> Context -> Options -> Baikai m [AssistantMessageEvent] + StreamCollectTraced :: TraceSink -> Model -> Context -> Options -> Baikai m [AssistantMessageEvent] StreamEach :: Model -> Context -> Options -> (AssistantMessageEvent -> m ()) -> Baikai m () ``` @@ -26,10 +27,40 @@ | --------------------------------- | ------------------------------------------------------------- | | `complete m c o` | Blocking completion → `Response`, including error-shaped responses. | | `streamCollect m c o` | Drain the stream into the full `[AssistantMessageEvent]`. | +| `streamCollectTraced sink m c o` | Collect events with tracing and bounded cleanup on cancellation. | | `streamEach m c o k` | Run callback `k` once per event, in order, inside `Eff`. | `streamEach` preserves incrementality (the callback sees events as they arrive); a terminal `EventError` appears in-band rather than as a thrown exception, matching baikai. + +## Evidence on cancellation + +`streamCollectTraced` delegates to core's `withTraceListWith`, which owns the +whole fold lifetime. Supply a `TraceSink` and opt in through `Options.evidence`. +On success or provider failure the sink's single `CallEvidence` carries exactly +the terminal's evidence record, including its `callId`. A strict pre-dispatch +refusal retains core's failed evidence and never invokes the provider. Opt-out +produces no evidence; providers that omit successful evidence retain core's +strict missing-record failure. + +For cancellation before terminal commitment, use a short, nonblocking sink that stores the complete +`ModelCallEvidence` in a caller-owned cell before acknowledging it. Read that +cell immediately in the caller's exception handler. With a healthy sink the +`CallAborted` record is already present and the original exception is unchanged. +The sink fold's end action also completes before propagation. Cancellation after +terminal commitment retains that terminal's evidence and adds no abort. + +The sink drain budget is one second in total. A blocked or throwing sink still +allows cancellation to unwind and reports its failure on stderr; an empty cell +means delivery was unconfirmed. A worker abandoned at the bound can deliver +later. Exactly-once means Baikai submits one evidence event, not that a failing +external destination durably stores it. Keep slow export behind a nonblocking +queue. + +Core's `withTraceFold`, `withTraceList`, `withTrace`, and their `With` variants +own the same lifetime. Arbitrary abandonment of `withTraceStream` between stream +steps still relies on garbage collection. The existing `streamCollect` and +`streamEach` operations keep their behavior. ## Interpreters diff --git a/baikai-effectful.cabal b/baikai-effectful.cabal --- a/baikai-effectful.cabal +++ b/baikai-effectful.cabal @@ -1,10 +1,10 @@ cabal-version: 3.4 name: baikai-effectful -version: 0.4.0.5 +version: 0.5.0.0 synopsis: effectful binding for the baikai AI-provider transport description: A thin, policy-free effectful binding over baikai's transport. Provides the dynamic - `Baikai` effect (Complete / StreamCollect / StreamEach) and interpreters over a real or + `Baikai` effect (Complete / StreamCollect / StreamCollectTraced / StreamEach) and interpreters over a real or isolated provider registry. Adds no retries, caching, budgets, or error remapping. category: AI @@ -60,7 +60,7 @@ hs-source-dirs: src exposed-modules: Baikai.Effectful build-depends: - baikai ^>=0.10.0, + baikai ^>=0.10.1, base >=4.20 && <5, effectful-core ^>=2.7, streamly-core >=0.3 && <0.5, @@ -81,6 +81,7 @@ LiveSpec StreamSpec StubProvider + TracedStreamSpec build-depends: baikai, @@ -92,4 +93,5 @@ tasty, tasty-hunit, text, + time, vector, diff --git a/src/Baikai/Effectful.hs b/src/Baikai/Effectful.hs --- a/src/Baikai/Effectful.hs +++ b/src/Baikai/Effectful.hs @@ -11,7 +11,7 @@ -- failures propagate exactly as baikai produces them — the blocking path -- ('complete') returns an error-shaped 'Response' whose -- 'Baikai.responseError' is populated; the streaming paths --- ('streamCollect', 'streamEach') surface baikai's terminal @EventError@ in-band. +-- ('streamCollect', 'streamCollectTraced', 'streamEach') surface baikai's terminal @EventError@ in-band. -- Retries, caching, budgets, rate limiting, and error remapping belong one layer -- up, in terms of this effect. module Baikai.Effectful @@ -21,6 +21,7 @@ -- * Operations complete, streamCollect, + streamCollectTraced, streamEach, -- * Interpreters @@ -33,12 +34,14 @@ Options, Response, AssistantMessageEvent, + TraceSink, ) where import Baikai (AssistantMessageEvent, Context, Model, Options, Response) import Baikai qualified import Baikai.Provider.Registry (ProviderRegistry) +import Baikai.Trace (TraceSink, withTraceListWith) import Effectful (Dispatch (Dynamic), DispatchOf, Eff, Effect, IOE, liftIO, (:>)) import Effectful.Dispatch.Dynamic (interpret, localSeqUnliftIO, send) import Streamly.Data.Fold qualified as Fold @@ -54,6 +57,8 @@ -- 'Baikai.streamRequest' drained to a list; a terminal @EventError@ appears -- as the last element on failure (baikai does not throw here). StreamCollect :: Model -> Context -> Options -> Baikai m [AssistantMessageEvent] + -- | Collect with fold-owned tracing and bounded cleanup on cancellation. + StreamCollectTraced :: TraceSink -> Model -> Context -> Options -> Baikai m [AssistantMessageEvent] -- | A streaming completion where each event is handed, in order, to a -- caller-supplied callback that runs inside the same effect context. StreamEach :: Model -> Context -> Options -> (AssistantMessageEvent -> m ()) -> Baikai m () @@ -71,6 +76,19 @@ streamCollect :: (Baikai :> es) => Model -> Context -> Options -> Eff es [AssistantMessageEvent] streamCollect m c o = send (StreamCollect m c o) +-- | Collect streamed events while sending call traces to a sink. When evidence +-- is requested and cancellation precedes terminal commitment, a healthy sink +-- receives Baikai's abort record before the original exception propagates +-- unchanged. After commitment the terminal evidence is retained. Sink cleanup waits at +-- most one second; a blocked or throwing sink leaves delivery unconfirmed and +-- reports on stderr. An abandoned worker may deliver later. Use a short sink +-- that stores the complete evidence in a caller-owned cell before acknowledging +-- it, then read that cell in the exception handler. Opt-out emits no evidence; +-- successful evidence is the provider terminal's own record, never reconstructed. +streamCollectTraced :: + (Baikai :> es) => TraceSink -> Model -> Context -> Options -> Eff es [AssistantMessageEvent] +streamCollectTraced sink m c o = send (StreamCollectTraced sink m c o) + -- | Stream a completion, invoking the callback on each event in order, inside -- 'Eff'. Preserves incrementality: the callback sees events as they arrive, in -- the same effect context as the caller. @@ -92,6 +110,8 @@ liftIO (Baikai.completeRequestWith reg m c o) StreamCollect m c o -> liftIO (Stream.toList (Baikai.streamRequestWith reg m c o)) + StreamCollectTraced sink m c o -> + liftIO (withTraceListWith reg sink m c o) StreamEach m c o k -> -- @k@ is the caller's @Eff@ action (a higher-order argument). @localSeqUnliftIO@ -- gives a function @unlift :: forall r. Eff localEs r -> IO r@ valid for the diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -6,6 +6,7 @@ import LiveSpec qualified import StreamSpec qualified import Test.Tasty (defaultMain, testGroup) +import TracedStreamSpec qualified main :: IO () main = @@ -14,5 +15,6 @@ "baikai-effectful" [ CompleteSpec.tests, StreamSpec.tests, + TracedStreamSpec.tests, LiveSpec.tests ] diff --git a/test/StubProvider.hs b/test/StubProvider.hs --- a/test/StubProvider.hs +++ b/test/StubProvider.hs @@ -10,6 +10,8 @@ -- agreement test depends on. module StubProvider ( stubRegistry, + evidenceRegistry, + blockingRegistry, stubModel, stubContext, stubOptions, @@ -18,7 +20,14 @@ where import Baikai +import Baikai.Evidence.Build qualified as Build import Baikai.Prelude +import Control.Concurrent (threadDelay) +import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar) +import Control.Exception (finally) +import Control.Monad (forever) +import Data.IORef (modifyIORef', newIORef, readIORef) +import Data.Time (getCurrentTime) import Data.Vector qualified as V import Streamly.Data.Stream (Stream) import Streamly.Data.Stream qualified as Stream @@ -103,3 +112,53 @@ (stubComplete t) ) pure reg + +-- | An evidence-aware fixture. Counts actual dispatches, so refusal tests can +-- distinguish a core refusal from an error emitted by the provider. +evidenceRegistry :: Bool -> IO (ProviderRegistry, IO Int) +evidenceRegistry failCall = do + counter <- newIORef 0 + reg <- newProviderRegistry + let handler m _ opts = do + modifyIORef' counter (+ 1) + now <- getCurrentTime + let err = if failCall then Just (providerError "fixture failure") else Nothing + ev <- + Build.minimalEvidence + m + opts + (Build.transportForModel m) + (Build.requestedTranslation opts) + (Build.dispatchEnvelope m opts) + now + now + (if failCall then CallFailed else CallSucceeded) + err + pure + ( stubResponse "evidence" + & #evidence + .~ ev + & #errorInfo + .~ err + & #message + . #stopReason + .~ (if failCall then ErrorReason else Stop) + & #message + . #errorMessage + .~ fmap (^. #message) err + ) + registerApiProviderWith reg (apiProviderWith stubApi (liftCompleteToStream handler) handler) + pure (reg, readIORef counter) + +-- | Block inside provider execution until cancelled; the cleanup barrier proves +-- the provider was released before the caller handles its exception. +blockingRegistry :: IO (ProviderRegistry, MVar (), MVar ()) +blockingRegistry = do + entered <- newEmptyMVar + released <- newEmptyMVar + reg <- newProviderRegistry + let handler _ _ _ = do + putMVar entered () + (forever (threadDelay 1000000)) `finally` putMVar released () + registerApiProviderWith reg (apiProviderWith stubApi (liftCompleteToStream handler) handler) + pure (reg, entered, released) diff --git a/test/TracedStreamSpec.hs b/test/TracedStreamSpec.hs new file mode 100644 --- /dev/null +++ b/test/TracedStreamSpec.hs @@ -0,0 +1,162 @@ +module TracedStreamSpec (tests) where + +import Baikai hiding (CallFailed, requested) +import Baikai.Effectful (runBaikai, runBaikaiWith, streamCollectTraced) +import Baikai.Evidence qualified as Ev +import Baikai.Prelude +import Baikai.Trace (TraceEvent (..), TraceSink (..)) +import Control.Concurrent (forkIO, throwTo) +import Control.Concurrent.MVar (newEmptyMVar, putMVar, readMVar, takeMVar, tryPutMVar, tryReadMVar) +import Control.Exception (Exception (..), SomeException, asyncExceptionFromException, asyncExceptionToException, finally, fromException, throwIO, try) +import Control.Monad (void, when) +import Data.IORef (IORef, modifyIORef', newIORef, readIORef) +import Data.Maybe (isJust) +import Effectful (runEff) +import GHC.Clock (getMonotonicTimeNSec) +import Streamly.Data.Fold qualified as Fold +import StubProvider +import System.Timeout (timeout) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +newtype Cancelled = Cancelled Int + deriving stock (Eq, Show) + +instance Exception Cancelled where + toException = asyncExceptionToException + fromException = asyncExceptionFromException + +requested :: Options +requested = stubOptions & #evidence .~ Just (evidenceRequest "effectful-abort") + +strictOptions :: Options +strictOptions = requested & #evidence . _Just . #strictness .~ EvidenceRequired EvidenceRequestedOnly + +memorySink :: IO (IORef [TraceEvent], TraceSink) +memorySink = do + cell <- newIORef [] + pure (cell, TraceSink (Fold.drainMapM (\e -> modifyIORef' cell (e :)))) + +records :: [TraceEvent] -> [ModelCallEvidence] +records events = [ev | CallEvidence {evidence = ev} <- events] + +tests :: TestTree +tests = + testGroup + "TracedStreamSpec" + [ terminalTest False, + terminalTest True, + testCase "strict refusal precedes dispatch and preserves terminal evidence" $ do + (reg, count) <- evidenceRegistry False + (cell, sink) <- memorySink + let opts = requested & #evidence . _Just . #strictness .~ EvidenceRequired EvidenceCorrelated + events <- runEff . runBaikaiWith reg $ streamCollectTraced sink stubModel stubContext opts + count >>= (@?= 0) + traces <- records <$> readIORef cell + let terminalRecords = [ev | EventError p <- events, Just ev <- [p ^. #evidence]] + traces @?= terminalRecords + length traces @?= 1 + map (^. #status) traces @?= [Ev.CallFailed] + assertBool "refusal error recorded" (all (isJust . (^. #errorInfo)) traces), + testCase "opt-out emits no evidence" $ do + (reg, _) <- evidenceRegistry False + (cell, sink) <- memorySink + _ <- runEff . runBaikaiWith reg $ streamCollectTraced sink stubModel stubContext stubOptions + readIORef cell >>= ((@?= []) . records), + testCase "a provider without evidence keeps strict missing-record enforcement" $ do + reg <- stubRegistry "legacy" + (cell, sink) <- memorySink + events <- runEff . runBaikaiWith reg $ streamCollectTraced sink stubModel stubContext strictOptions + assertBool "missing record fails" (any (\case EventError {} -> True; _ -> False) events) + readIORef cell >>= ((@?= []) . records), + testCase "global interpreter delegates traced collection" $ do + let model = stubModel & #api .~ Custom "effectful-traced-global-102" + handler _ _ _ = pure emptyResponse + registerApiProvider (apiProviderWith (model ^. #api) (liftCompleteToStream handler) handler) + (cell, sink) <- memorySink + events <- runEff . runBaikai $ streamCollectTraced sink model stubContext stubOptions + assertBool "terminal returned" (any (\case EventDone {} -> True; _ -> False) events) + traces <- readIORef cell + length traces @?= 2, + testCase "best-effort cancellation acknowledges abort before handler" $ cancellationTest requested Healthy, + testCase "strict cancellation preserves the original exception" $ cancellationTest strictOptions Healthy, + testCase "blocked sink bounds cancellation and may deliver later once" $ cancellationTest strictOptions Blocked, + testCase "blocked sink also bounds best-effort cancellation" $ cancellationTest requested Blocked, + testCase "throwing sink cannot replace strict cancellation" $ cancellationTest strictOptions Throwing, + testCase "throwing sink cannot replace best-effort cancellation" $ cancellationTest requested Throwing + ] + +terminalTest :: Bool -> TestTree +terminalTest failCall = testCase (if failCall then "provider failure keeps evidence identity" else "success keeps evidence identity") $ do + (reg, count) <- evidenceRegistry failCall + (cell, sink) <- memorySink + events <- runEff . runBaikaiWith reg $ streamCollectTraced sink stubModel stubContext requested + count >>= (@?= 1) + traces <- records <$> readIORef cell + let terminalRecords = + [ ev + | event <- events, + p <- case event of EventDone t -> [t]; EventError t -> [t]; _ -> [], + Just ev <- [p ^. #evidence] + ] + traces @?= terminalRecords + length traces @?= 1 + map (^. #status) traces @?= [if failCall then Ev.CallFailed else CallSucceeded] + +data SinkMode = Healthy | Blocked | Throwing + deriving stock (Eq) + +cancellationTest :: Options -> SinkMode -> IO () +cancellationTest opts mode = do + (reg, entered, released) <- blockingRegistry + cell <- newIORef [] + sinkEntered <- newEmptyMVar + sinkReleased <- newEmptyMVar + ended <- newEmptyMVar + let step e = do + case e of + CallStarted {} -> do + putMVar sinkEntered () + case mode of + Blocked -> readMVar sinkReleased + Throwing -> throwIO (userError "fixture sink failure") + Healthy -> pure () + _ -> pure () + modifyIORef' cell (e :) + sink = TraceSink (Fold.rmapM (\() -> putMVar ended ()) (Fold.drainMapM step)) + outcome <- newEmptyMVar + tid <- forkIO $ do + result <- try (runEff . runBaikaiWith reg $ streamCollectTraced sink stubModel stubContext opts) + -- All observations belong to the caller's exception handler: no waiting + -- for evidence, forced collection, or post-cancellation polling. + snapshot <- readIORef cell + providerDone <- tryReadMVar released + sinkDone <- tryReadMVar ended + putMVar outcome (result :: Either SomeException [AssistantMessageEvent], snapshot, providerDone, sinkDone) + let run = do + timeout 2000000 (readMVar entered) >>= (@?= Just ()) + timeout 2000000 (readMVar sinkEntered) >>= (@?= Just ()) + start <- getMonotonicTimeNSec + throwTo tid (Cancelled 73102) + observed <- timeout 3000000 (takeMVar outcome) + stop <- getMonotonicTimeNSec + case observed of + Nothing -> assertFailure "cancelled collection hung" + Just (result, snapshot, providerDone, sinkDone) -> do + case result of + Left ex -> fromException ex @?= Just (Cancelled 73102) + Right _ -> assertFailure "cancellation disappeared" + providerDone @?= Just () + case mode of + Healthy -> do + map (^. #status) (records snapshot) @?= [Ev.CallAborted] + sinkDone @?= Just () + _ -> records snapshot @?= [] + assertBool "cleanup stays bounded" (stop - start < 2500000000) + when (mode == Blocked) $ do + assertBool "production drain budget was exercised" (stop - start >= 800000000) + putMVar sinkReleased () + timeout 2000000 (readMVar ended) >>= (@?= Just ()) + eventual <- records <$> readIORef cell + map (^. #status) eventual @?= [Ev.CallAborted] + run `finally` (void (tryPutMVar sinkReleased ()) >> throwTo tid (Cancelled 73102))