diff --git a/baikai-trace-otel.cabal b/baikai-trace-otel.cabal --- a/baikai-trace-otel.cabal +++ b/baikai-trace-otel.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: baikai-trace-otel -version: 0.2.0.0 +version: 0.3.0.0 synopsis: OpenTelemetry TraceSink for baikai. description: Provides an opt-in OpenTelemetry adapter for the baikai 'TraceSink' @@ -35,7 +35,7 @@ hs-source-dirs: src exposed-modules: Baikai.Trace.Sink.OpenTelemetry build-depends: - , baikai ^>=0.2.0 + , baikai ^>=0.3.0 , base >=4.20 && <5 , containers , hs-opentelemetry-api >=1.0 && <1.1 @@ -53,7 +53,7 @@ main-is: Main.hs ghc-options: -threaded -with-rtsopts=-N build-depends: - , baikai ^>=0.2.0 + , baikai ^>=0.3.0 , baikai-trace-otel , base , generic-lens @@ -61,6 +61,7 @@ , hs-opentelemetry-exporter-in-memory >=1.0 && <1.1 , hs-opentelemetry-sdk >=1.0 && <1.1 , lens ^>=5.3 + , streamly-core ^>=0.3 , tasty , tasty-hunit , text diff --git a/src/Baikai/Trace/Sink/OpenTelemetry.hs b/src/Baikai/Trace/Sink/OpenTelemetry.hs --- a/src/Baikai/Trace/Sink/OpenTelemetry.hs +++ b/src/Baikai/Trace/Sink/OpenTelemetry.hs @@ -107,6 +107,9 @@ pure (Map.insert eventId sp m) CallFinished {eventId, timestamp, model, latencyMs, inputTokens, outputTokens, usd} -> case Map.lookup eventId m of + -- A terminal without a live span is unreachable for normal withTraceStream + -- usage because each traced call drives a fresh fold. Keep the silent drop so + -- hand-fed or replayed event streams cannot crash the sink. Nothing -> pure m Just sp -> do let attrs :: HashMap.HashMap Text Attr.Attribute @@ -116,7 +119,7 @@ maybe id (\s -> HashMap.insert "baikai.cost.usd" (Attr.toAttribute (Scientific.toRealFloat s :: Double))) usd $ AttrMap.insertByKey SC.genAi_response_model model $ HashMap.fromList - [ ("baikai.latency_ms", Attr.toAttribute (fromIntegral latencyMs :: Int)) + [ ("baikai.latency_ms", Attr.toAttribute latencyMs) ] Otel.addAttributes sp attrs Otel.setStatus sp Otel.Ok @@ -124,11 +127,14 @@ pure (Map.delete eventId m) CallFailed {eventId, timestamp, latencyMs, errorMessage} -> case Map.lookup eventId m of + -- A terminal without a live span is unreachable for normal withTraceStream + -- usage because each traced call drives a fresh fold. Keep the silent drop so + -- hand-fed or replayed event streams cannot crash the sink. Nothing -> pure m Just sp -> do Otel.addAttributes sp $ HashMap.fromList - [ ("baikai.latency_ms", Attr.toAttribute (fromIntegral latencyMs :: Int)), + [ ("baikai.latency_ms", Attr.toAttribute latencyMs), ("baikai.error", Attr.toAttribute errorMessage) ] Otel.setStatus sp (Otel.Error errorMessage) diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -4,28 +4,32 @@ import Baikai.Api (Api (..)) import Baikai.Content (AssistantContent (..), TextContent (..)) -import Baikai.Context (Context (..), _Context) +import Baikai.Context (Context (..), emptyContext) import Baikai.Error (BaikaiError, providerError) import Baikai.Message (AssistantPayload (..), user) -import Baikai.Model (Model (..), _Model) -import Baikai.Options (Options, _Options) +import Baikai.Model (Model (..), emptyModel) +import Baikai.Options (Options, emptyOptions) import Baikai.Provider (ApiProvider (..), registerApiProvider) import Baikai.Response (Response (..)) import Baikai.StopReason (StopReason (..)) import Baikai.Stream (liftCompleteToStream) -import Baikai.Trace (withTrace) +import Baikai.Trace (withTrace, withTraceStream) import Baikai.Trace.Sink.OpenTelemetry (otelSink) -import Baikai.Usage (Usage, _Usage) +import Baikai.Usage (Usage, zeroUsage) +import Control.Concurrent (threadDelay) import Control.Exception (throwIO) import Control.Lens ((&), (.~), (^.)) import Data.Generics.Labels () import Data.HashMap.Strict qualified as HashMap import Data.IORef (IORef, readIORef) import Data.Text (Text) +import Data.Text qualified as Text import Data.Vector qualified as V import OpenTelemetry.Attributes qualified as Attr import OpenTelemetry.Exporter.InMemory.Span (inMemoryListExporter) import OpenTelemetry.Trace.Core qualified as Otel +import Streamly.Data.Stream qualified as Stream +import System.Mem (performMajorGC) import Test.Tasty (TestTree, defaultMain, testGroup) import Test.Tasty.HUnit (assertBool, assertEqual, assertFailure, testCase, (@?=)) @@ -35,7 +39,8 @@ testGroup "baikai-trace-otel" [ successSpanTest, - failureSpanTest + failureSpanTest, + abortSpanTest ] -- | Build a stub 'Model' under a private 'Api' tag. Each test uses @@ -43,21 +48,21 @@ -- registry between tests. stubModel :: Api -> Model stubModel a = - _Model + emptyModel & #modelId .~ "stub-1" & #api .~ a & #provider .~ "stub.otel" & #maxOutputTokens .~ 16 stubContext :: Context -stubContext = _Context & #messages .~ V.fromList [user "hello"] +stubContext = emptyContext & #messages .~ V.fromList [user "hello"] stubOptions :: Options -stubOptions = _Options & #maxTokens .~ Just 16 +stubOptions = emptyOptions & #maxTokens .~ Just 16 sampleUsage :: Usage sampleUsage = - _Usage + zeroUsage & #inputTokens .~ 12 & #outputTokens .~ 3 & #totalTokens .~ 15 @@ -71,7 +76,7 @@ usage = sampleUsage, stopReason = Stop, errorMessage = Nothing, - timestamp = read "2026-05-14 00:00:00 UTC" + timestamp = Just (read "2026-05-14 00:00:00 UTC") }, model = stubModel a, api = a, @@ -177,3 +182,41 @@ assertBool "has baikai.error" (HashMap.member "baikai.error" attrs) assertBool "has baikai.latency_ms" (HashMap.member "baikai.latency_ms" attrs) _ -> assertFailure "expected exactly one span" + +abortSpanTest :: TestTree +abortSpanTest = + testCase "early abort closes the span with Error status" $ do + let a = Custom "baikai-otel-abort" + registerOk a + (tracer, getSpans) <- newTracerWithInMemory + let sink = otelSink tracer + emitted <- + Stream.toList + (Stream.take 1 (withTraceStream sink (stubModel a) stubContext stubOptions)) + length emitted @?= 1 + spans <- awaitSpans getSpans 1 + assertEqual "exactly one span recorded" 1 (length spans) + case spans of + [sp] -> do + hot <- spanHotSnapshot sp + case Otel.hotStatus hot of + Otel.Error msg -> + assertBool + ("expected abort message, got: " <> show msg) + ("aborted" `Text.isInfixOf` msg) + other -> assertFailure ("expected Error status, got: " <> show other) + _ -> assertFailure "expected exactly one span" + +-- The trace finalizer on an abandoned stream runs from streamly's GC hook. +awaitSpans :: IO [Otel.ImmutableSpan] -> Int -> IO [Otel.ImmutableSpan] +awaitSpans getSpans n = go (100 :: Int) + where + go 0 = do + spans <- getSpans + assertFailure ("timed out waiting for spans; got: " <> show (length spans)) + go k = do + performMajorGC + spans <- getSpans + if length spans >= n + then pure spans + else threadDelay 50000 >> go (k - 1)