diff --git a/pkgs/sdk/server-ai/src/Config/LdAiConfig.cs b/pkgs/sdk/server-ai/src/Config/LdAiConfig.cs index 7fe18702..37221a5b 100644 --- a/pkgs/sdk/server-ai/src/Config/LdAiConfig.cs +++ b/pkgs/sdk/server-ai/src/Config/LdAiConfig.cs @@ -1,13 +1,15 @@ +using System; using System.Collections.Generic; using System.Linq; using LaunchDarkly.Sdk.Server.Ai.DataModel; +using LaunchDarkly.Sdk.Server.Ai.Interfaces; namespace LaunchDarkly.Sdk.Server.Ai.Config; /// /// Represents an AI Config, which contains model parameters and prompt messages. /// -public record LdAiConfig +public class LdAiConfig { /// @@ -217,7 +219,10 @@ public LdAiConfig Build() /// public readonly ModelProvider Provider; - internal LdAiConfig(bool enabled, IEnumerable messages, Meta meta, Model model, Provider provider) + private readonly Func _trackerFactory; + + internal LdAiConfig(bool enabled, IEnumerable messages, Meta meta, Model model, Provider provider, + Func createTracker = null) { Model = new ModelConfiguration(model?.Name ?? "", model?.Parameters ?? new Dictionary(), model?.Custom ?? new Dictionary()); @@ -226,7 +231,20 @@ internal LdAiConfig(bool enabled, IEnumerable messages, Meta meta, Mode Version = meta?.Version ?? 1; Enabled = enabled; Provider = new ModelProvider(provider?.Name ?? ""); + _trackerFactory = createTracker; + } + + internal LdAiConfig(LdAiConfig source, Func createTracker) + { + Model = source.Model; + Messages = source.Messages; + VariationKey = source.VariationKey; + Version = source.Version; + Enabled = source.Enabled; + Provider = source.Provider; + _trackerFactory = createTracker; } + internal LdValue ToLdValue() { return LdValue.ObjectFrom(new Dictionary @@ -278,6 +296,17 @@ internal LdValue ToLdValue() /// public int Version { get; } + /// + /// Returns a new tracker for a fresh AI run. Each call mints a new runId (a UUIDv4) + /// that LaunchDarkly uses to correlate the run's events in metrics views. Call this + /// once per AI run; metrics from different runIds cannot be combined. + /// + /// + /// Returns null for configs created directly via the builder; only configs returned by + /// have a factory wired in. + /// + public ILdAiConfigTracker CreateTracker() => _trackerFactory?.Invoke(); + /// /// Convenient helper that returns a disabled LdAiConfig. /// diff --git a/pkgs/sdk/server-ai/src/Interfaces/ILdAiClient.cs b/pkgs/sdk/server-ai/src/Interfaces/ILdAiClient.cs index 28fd3dfb..526c7cbd 100644 --- a/pkgs/sdk/server-ai/src/Interfaces/ILdAiClient.cs +++ b/pkgs/sdk/server-ai/src/Interfaces/ILdAiClient.cs @@ -12,8 +12,9 @@ public interface ILdAiClient /// /// Retrieves a LaunchDarkly AI Completion Config identified by the given key. The return value - /// is an , which makes the configuration available and - /// provides convenience methods for generating events related to model usage. + /// is an , which makes the configuration available and provides + /// a factory for generating trackers capable of + /// emitting events related to model usage. /// /// Any variables provided will be interpolated into the prompt's messages. /// Additionally, the current LaunchDarkly context will be available as 'ldctx' within @@ -25,8 +26,8 @@ public interface ILdAiClient /// the default config, if unable to retrieve from LaunchDarkly. When not provided, /// a disabled config is used as the fallback. /// the list of variables used when interpolating the prompt - /// an AI Completion Config tracker - public ILdAiConfigTracker CompletionConfig(string key, Context context, LdAiConfig defaultValue = null, + /// an AI Completion Config + public LdAiConfig CompletionConfig(string key, Context context, LdAiConfig defaultValue = null, IReadOnlyDictionary variables = null); /// @@ -36,8 +37,21 @@ public ILdAiConfigTracker CompletionConfig(string key, Context context, LdAiConf /// the context /// the default config, if unable to retrieve from LaunchDarkly /// the list of variables used when interpolating the prompt - /// an AI Completion Config tracker + /// an AI Completion Config [Obsolete("Use CompletionConfig instead.")] - public ILdAiConfigTracker Config(string key, Context context, LdAiConfig defaultValue = null, + public LdAiConfig Config(string key, Context context, LdAiConfig defaultValue = null, IReadOnlyDictionary variables = null); + + /// + /// Reconstructs a tracker from a resumption token. This enables cross-process scenarios + /// such as deferred feedback, where a tracker's runId needs to be reused in a different + /// process or at a later time. + /// + /// The reconstructed tracker will have empty model and provider names, as these are not + /// included in the resumption token. + /// + /// the resumption token obtained from + /// the context to use for track events + /// a tracker associated with the original runId + public ILdAiConfigTracker CreateTracker(string resumptionToken, Context context); } diff --git a/pkgs/sdk/server-ai/src/Interfaces/ILdAiConfigTracker.cs b/pkgs/sdk/server-ai/src/Interfaces/ILdAiConfigTracker.cs index 422bb261..b89e20fa 100644 --- a/pkgs/sdk/server-ai/src/Interfaces/ILdAiConfigTracker.cs +++ b/pkgs/sdk/server-ai/src/Interfaces/ILdAiConfigTracker.cs @@ -6,21 +6,36 @@ namespace LaunchDarkly.Sdk.Server.Ai.Interfaces; /// -/// A utility capable of generating events related to a specific AI model -/// configuration. +/// Records metrics for a single AI run. /// +/// +/// All events emitted by a tracker share a runId (a UUIDv4) so LaunchDarkly can correlate +/// them. See individual track methods for their specific semantics. +/// Call CreateTracker on the AI Config to start a new run. A +/// preserves the runId, so events emitted by a tracker +/// reconstructed in another process correlate with the original tracker's runId. +/// public interface ILdAiConfigTracker { /// - /// The AI model configuration retrieved from LaunchDarkly, or a default value if unable to retrieve. + /// A URL-safe Base64-encoded token that can be used to reconstruct this tracker in a different + /// process or at a later time. The token contains the runId, configKey, variationKey, and version. + /// + /// Use to reconstruct a tracker from this token. /// - public LdAiConfig Config { get; } + public string ResumptionToken { get; } + + /// + /// A summary of the metrics tracked by this tracker. + /// + public MetricSummary Summary { get; } /// /// Tracks a duration metric related to this config. For example, if a particular operation /// related to usage of the AI model takes 100ms, this can be tracked and made available in /// LaunchDarkly. /// + /// Records at most once per Tracker; further calls are ignored. /// the duration in milliseconds public void TrackDuration(float durationMs); @@ -36,16 +51,18 @@ public interface ILdAiConfigTracker /// type of the task's result /// the task public Task TrackDurationOfTask(Task task); - + /// /// Tracks the time it takes for the first token to be generated. /// + /// Records at most once per Tracker; further calls are ignored. /// the duration in milliseconds public void TrackTimeToFirstToken(float timeToFirstTokenMs); /// /// Tracks feedback (positive or negative) related to the output of the model. /// + /// Records at most once per Tracker; further calls are ignored. /// the feedback /// thrown if the feedback value is not or public void TrackFeedback(Feedback feedback); @@ -53,11 +70,19 @@ public interface ILdAiConfigTracker /// /// Tracks a generation event related to this config. /// + /// + /// Records at most once per Tracker. TrackSuccess and TrackError share state; only + /// one of the two can record per Tracker, and subsequent calls are ignored. + /// public void TrackSuccess(); /// /// Tracks an unsuccessful generation event related to this config. /// + /// + /// Records at most once per Tracker. TrackSuccess and TrackError share state; only + /// one of the two can record per Tracker, and subsequent calls are ignored. + /// public void TrackError(); /// @@ -92,6 +117,10 @@ public interface ILdAiConfigTracker /// Task is automatically measured and recorded as the latency metric associated with this request. /// /// + /// + /// Subsequent calls re-run the task but emit only metrics not already recorded on this Tracker. + /// Call CreateTracker on the AI Config to start a new run. + /// /// a task representing the request /// the task public Task TrackRequest(Task request); @@ -99,6 +128,7 @@ public interface ILdAiConfigTracker /// /// Tracks token usage related to this config. /// + /// Records at most once per Tracker; further calls are ignored. /// the token usage public void TrackTokens(Usage usage); } diff --git a/pkgs/sdk/server-ai/src/LdAiClient.cs b/pkgs/sdk/server-ai/src/LdAiClient.cs index 803b3d5a..18e70457 100644 --- a/pkgs/sdk/server-ai/src/LdAiClient.cs +++ b/pkgs/sdk/server-ai/src/LdAiClient.cs @@ -58,7 +58,7 @@ public LdAiClient(ILaunchDarklyClient client) private const string LdContextVariable = "ldctx"; /// - public ILdAiConfigTracker CompletionConfig(string key, Context context, LdAiConfig defaultValue = null, + public LdAiConfig CompletionConfig(string key, Context context, LdAiConfig defaultValue = null, IReadOnlyDictionary variables = null) { _client.Track(TrackUsageCompletionConfig, context, LdValue.Of(key), 1); @@ -72,7 +72,7 @@ public ILdAiConfigTracker CompletionConfig(string key, Context context, LdAiConf /// This allows higher-level SDK entry methods to track their own usage events without /// double-counting. /// - private ILdAiConfigTracker Evaluate(string key, Context context, LdAiConfig defaultValue, + private LdAiConfig Evaluate(string key, Context context, LdAiConfig defaultValue, IReadOnlyDictionary variables = null) { var result = _client.JsonVariation(key, context, defaultValue.ToLdValue()); @@ -80,8 +80,7 @@ private ILdAiConfigTracker Evaluate(string key, Context context, LdAiConfig defa var parsed = ParseConfig(result, key); if (parsed == null) { - // ParseConfig already does logging. - return new LdAiConfigTracker(_client, key, defaultValue, context); + return WithTrackerFactory(defaultValue, key, context); } var mergedVariables = new Dictionary { { LdContextVariable, GetAllAttributes(context) } }; @@ -113,12 +112,28 @@ private ILdAiConfigTracker Evaluate(string key, Context context, LdAiConfig defa { _logger.Error( $"AI model config prompt has malformed message at index {i}: {ex.Message} (returning default config, which will not contain interpolated prompt messages)"); - return new LdAiConfigTracker(_client, key, defaultValue, context); + return WithTrackerFactory(defaultValue, key, context); } } } - return new LdAiConfigTracker(_client, key, new LdAiConfig(parsed.Meta?.Enabled ?? false, prompt, parsed.Meta, parsed.Model, parsed.Provider), context); + var config = new LdAiConfig(parsed.Meta?.Enabled ?? false, prompt, parsed.Meta, parsed.Model, parsed.Provider); + return WithTrackerFactory(config, key, context); + } + + private LdAiConfig WithTrackerFactory(LdAiConfig config, string configKey, Context context) + { + var variationKey = config.VariationKey; + var version = config.Version; + var modelName = config.Model?.Name ?? ""; + var providerName = config.Provider?.Name ?? ""; + var client = _client; + + return new LdAiConfig(config, () => + { + var runId = Guid.NewGuid().ToString(); + return new LdAiConfigTracker(client, runId, configKey, variationKey, version, context, modelName, providerName); + }); } /// @@ -128,15 +143,22 @@ private ILdAiConfigTracker Evaluate(string key, Context context, LdAiConfig defa /// the context /// the default config, if unable to retrieve from LaunchDarkly /// the list of variables used when interpolating the prompt - /// an AI Completion Config tracker + /// an AI Completion Config [Obsolete("Use CompletionConfig instead.")] - public ILdAiConfigTracker Config(string key, Context context, LdAiConfig defaultValue = null, + public LdAiConfig Config(string key, Context context, LdAiConfig defaultValue = null, IReadOnlyDictionary variables = null) { return CompletionConfig(key, context, defaultValue, variables); } + /// + public ILdAiConfigTracker CreateTracker(string resumptionToken, Context context) + { + return LdAiConfigTracker.FromResumptionToken(resumptionToken, _client, context); + } + + private static IDictionary AddSingleKindContextAttributes(Context context) { var attributes = new Dictionary diff --git a/pkgs/sdk/server-ai/src/LdAiConfigTracker.cs b/pkgs/sdk/server-ai/src/LdAiConfigTracker.cs index 93bfd372..266ef581 100644 --- a/pkgs/sdk/server-ai/src/LdAiConfigTracker.cs +++ b/pkgs/sdk/server-ai/src/LdAiConfigTracker.cs @@ -1,7 +1,12 @@ using System; using System.Collections.Generic; using System.Diagnostics; +using System.IO; using System.Runtime.CompilerServices; +using System.Text; +using System.Text.Json; +using System.Text.Json.Serialization; +using System.Threading; using System.Threading.Tasks; using LaunchDarkly.Sdk.Server.Ai.Config; using LaunchDarkly.Sdk.Server.Ai.Interfaces; @@ -10,13 +15,33 @@ namespace LaunchDarkly.Sdk.Server.Ai; /// -/// A tracker capable of reporting events related to a particular AI Config. +/// Records metrics for a single AI run. /// +/// +/// All events a tracker emits share a runId (a UUIDv4) so LaunchDarkly can correlate +/// them in metrics views. See individual track methods for their specific semantics. +/// Call CreateTracker on the AI Config to start a new run. A +/// preserves the runId, so events emitted by a tracker +/// reconstructed in another process correlate with the original tracker's runId. +/// public class LdAiConfigTracker : ILdAiConfigTracker { private readonly ILaunchDarklyClient _client; + private readonly string _runId; + private readonly string _configKey; + private readonly string _variationKey; + private readonly int _version; private readonly Context _context; + private readonly string _modelName; + private readonly string _providerName; private readonly LdValue _trackData; + private readonly ILogger _logger; + + private StrongBox _durationMs; + private StrongBox _timeToFirstTokenMs; + private StrongBox _tokens; + private StrongBox _feedback; + private StrongBox _trackedSuccess; // true = success, false = error private const string Duration = "$ld:ai:duration:total"; private const string FeedbackPositive = "$ld:ai:feedback:user:positive"; @@ -29,36 +54,88 @@ public class LdAiConfigTracker : ILdAiConfigTracker private const string TimeToFirstToken = "$ld:ai:tokens:ttf"; /// - /// Constructs a new AI Config tracker. The tracker is associated with a configuration, - /// a context, and a key which identifies the configuration. + /// Constructs a new AI Config tracker. The tracker is associated with a run, a config key, + /// and a context. The runId should be a unique identifier (UUID v4) for each AI run. /// /// the LaunchDarkly client + /// the unique run identifier /// key of the AI Config - /// the AI Config + /// the variation key + /// the config version /// the context - /// - public LdAiConfigTracker(ILaunchDarklyClient client, string configKey, LdAiConfig config, Context context) + /// the model name + /// the provider name + public LdAiConfigTracker(ILaunchDarklyClient client, string runId, string configKey, + string variationKey, int version, Context context, string modelName, string providerName) { - Config = config ?? throw new ArgumentNullException(nameof(config)); - _client = client ?? throw new ArgumentNullException(nameof(client)); + _client = client; + _runId = runId ?? ""; + _configKey = configKey ?? ""; + _variationKey = variationKey; + _version = version; _context = context; - _trackData = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey)}, - { "version", LdValue.Of(config.Version)}, - { "configKey" , LdValue.Of(configKey ?? throw new ArgumentNullException(nameof(configKey))) }, - { "modelName", LdValue.Of(config.Model?.Name) }, - { "providerName", LdValue.Of(config.Provider?.Name) }, - }); + _modelName = modelName ?? ""; + _providerName = providerName ?? ""; + _logger = client?.GetLogger(); + + var trackDataBuilder = new Dictionary + { + { "runId", LdValue.Of(_runId) }, + { "configKey", LdValue.Of(_configKey) }, + { "version", LdValue.Of(_version) }, + { "modelName", LdValue.Of(_modelName) }, + { "providerName", LdValue.Of(_providerName) }, + }; + if (!string.IsNullOrEmpty(_variationKey)) + { + trackDataBuilder.Add("variationKey", LdValue.Of(_variationKey)); + } + _trackData = LdValue.ObjectFrom(trackDataBuilder); } + /// + public string ResumptionToken + { + get + { + using var stream = new MemoryStream(); + using (var writer = new Utf8JsonWriter(stream)) + { + writer.WriteStartObject(); + writer.WriteString("runId", _runId); + writer.WriteString("configKey", _configKey); + if (!string.IsNullOrEmpty(_variationKey)) + { + writer.WriteString("variationKey", _variationKey); + } + writer.WriteNumber("version", _version); + writer.WriteEndObject(); + } + var base64 = Convert.ToBase64String(stream.ToArray()); + return base64.Replace('+', '-').Replace('/', '_').TrimEnd('='); + } + } /// - public LdAiConfig Config { get; } + public MetricSummary Summary => new MetricSummary( + _durationMs?.Value, + _feedback?.Value, + _tokens?.Value, + _trackedSuccess?.Value, + _timeToFirstTokenMs?.Value + ); /// - public void TrackDuration(float durationMs) => + public void TrackDuration(float durationMs) + { + if (Interlocked.CompareExchange(ref _durationMs, + new StrongBox(durationMs), null) != null) + { + _logger?.Warn("Skipping TrackDuration: duration already recorded on this tracker. Call CreateTracker on the AI Config for a new run. {0}", _trackData.ToJsonString()); + return; + } _client.Track(Duration, _context, _trackData, durationMs); + } /// @@ -74,34 +151,63 @@ public async Task TrackDurationOfTask(Task task) } /// - public void TrackTimeToFirstToken(float timeToFirstTokenMs) => + public void TrackTimeToFirstToken(float timeToFirstTokenMs) + { + if (Interlocked.CompareExchange(ref _timeToFirstTokenMs, + new StrongBox(timeToFirstTokenMs), null) != null) + { + _logger?.Warn("Skipping TrackTimeToFirstToken: time-to-first-token already recorded on this tracker. Call CreateTracker on the AI Config for a new run. {0}", _trackData.ToJsonString()); + return; + } _client.Track(TimeToFirstToken, _context, _trackData, timeToFirstTokenMs); + } /// public void TrackFeedback(Feedback feedback) { + // Validate the enum value first so invalid input throws without consuming the slot. + string eventName; switch (feedback) { case Feedback.Positive: - _client.Track(FeedbackPositive, _context, _trackData, 1); + eventName = FeedbackPositive; break; case Feedback.Negative: - _client.Track(FeedbackNegative, _context, _trackData, 1); + eventName = FeedbackNegative; break; default: throw new ArgumentOutOfRangeException(nameof(feedback), feedback, null); } + if (Interlocked.CompareExchange(ref _feedback, + new StrongBox(feedback), null) != null) + { + _logger?.Warn("Skipping TrackFeedback: feedback already recorded on this tracker. Call CreateTracker on the AI Config for a new run. {0}", _trackData.ToJsonString()); + return; + } + _client.Track(eventName, _context, _trackData, 1); } /// public void TrackSuccess() { + if (Interlocked.CompareExchange(ref _trackedSuccess, + new StrongBox(true), null) != null) + { + _logger?.Warn("Skipping TrackSuccess: success/error already recorded on this tracker. Call CreateTracker on the AI Config for a new run. {0}", _trackData.ToJsonString()); + return; + } _client.Track(GenerationSuccess, _context, _trackData, 1); } /// public void TrackError() { + if (Interlocked.CompareExchange(ref _trackedSuccess, + new StrongBox(false), null) != null) + { + _logger?.Warn("Skipping TrackError: success/error already recorded on this tracker. Call CreateTracker on the AI Config for a new run. {0}", _trackData.ToJsonString()); + return; + } _client.Track(GenerationError, _context, _trackData, 1); } @@ -137,6 +243,18 @@ public async Task TrackRequest(Task request) /// public void TrackTokens(Usage usage) { + // Empty usage doesn't burn the slot. + if ((usage.Total ?? 0) <= 0 && (usage.Input ?? 0) <= 0 && (usage.Output ?? 0) <= 0) + { + return; + } + // Atomic claim. + if (Interlocked.CompareExchange(ref _tokens, + new StrongBox(usage), null) != null) + { + _logger?.Warn("Skipping TrackTokens: tokens already recorded on this tracker. Call CreateTracker on the AI Config for a new run. {0}", _trackData.ToJsonString()); + return; + } if (usage.Total is > 0) { _client.Track(TokenTotal, _context, _trackData, usage.Total.Value); @@ -150,4 +268,66 @@ public void TrackTokens(Usage usage) _client.Track(TokenOutput, _context, _trackData, usage.Output.Value); } } + + /// + /// Reconstructs a tracker from a resumption token. This enables cross-process scenarios + /// such as deferred feedback, where a tracker's runId needs to be reused in a different + /// process or at a later time. + /// + /// The reconstructed tracker will have empty model and provider names, as these are not + /// included in the resumption token. + /// + /// the resumption token obtained from + /// the LaunchDarkly client + /// the context to use for track events + /// a new tracker associated with the original tracker's runId + /// thrown if token or client is null + /// thrown if the token is malformed or missing required fields + public static LdAiConfigTracker FromResumptionToken(string token, ILaunchDarklyClient client, Context context) + { + if (token == null) throw new ArgumentNullException(nameof(token)); + if (client == null) throw new ArgumentNullException(nameof(client)); + + var base64 = token.Replace('-', '+').Replace('_', '/'); + switch (base64.Length % 4) + { + case 2: base64 += "=="; break; + case 3: base64 += "="; break; + } + + ResumptionPayload payload; + try + { + var json = Encoding.UTF8.GetString(Convert.FromBase64String(base64)); + payload = JsonSerializer.Deserialize(json); + } + catch (Exception e) when (e is FormatException || e is JsonException) + { + throw new ArgumentException("Invalid resumption token", nameof(token), e); + } + + if (payload == null || string.IsNullOrEmpty(payload.RunId) || string.IsNullOrEmpty(payload.ConfigKey)) + { + throw new ArgumentException("Resumption token is missing required fields (runId, configKey)", + nameof(token)); + } + + return new LdAiConfigTracker(client, payload.RunId, payload.ConfigKey, + payload.VariationKey, payload.Version, context, "", ""); + } + + private class ResumptionPayload + { + [JsonPropertyName("runId")] + public string RunId { get; set; } + + [JsonPropertyName("configKey")] + public string ConfigKey { get; set; } + + [JsonPropertyName("variationKey")] + public string VariationKey { get; set; } + + [JsonPropertyName("version")] + public int Version { get; set; } + } } diff --git a/pkgs/sdk/server-ai/src/Tracking/MetricSummary.cs b/pkgs/sdk/server-ai/src/Tracking/MetricSummary.cs new file mode 100644 index 00000000..bfe79cf2 --- /dev/null +++ b/pkgs/sdk/server-ai/src/Tracking/MetricSummary.cs @@ -0,0 +1,17 @@ +namespace LaunchDarkly.Sdk.Server.Ai.Tracking; + +/// +/// A summary of the metrics tracked by a tracker. +/// +/// the duration in milliseconds +/// the feedback sentiment +/// the token usage +/// whether the generation was successful +/// the time to first token in milliseconds +public record struct MetricSummary( + double? DurationMs, + Feedback? Feedback, + Usage? Tokens, + bool? Success, + double? TimeToFirstTokenMs +); diff --git a/pkgs/sdk/server-ai/test/InterpolationTests.cs b/pkgs/sdk/server-ai/test/InterpolationTests.cs index 366f093e..a9f724d5 100644 --- a/pkgs/sdk/server-ai/test/InterpolationTests.cs +++ b/pkgs/sdk/server-ai/test/InterpolationTests.cs @@ -39,9 +39,9 @@ private string Eval(string prompt, Context context, IReadOnlyDictionary x.GetLogger()).Returns(mockLogger.Object); var client = new LdAiClient(mockClient.Object); - var tracker = client.Config("foo", context, LdAiConfig.Disabled, variables); + var config = client.Config("foo", context, LdAiConfig.Disabled, variables); - return tracker.Config.Messages[0].Content; + return config.Messages[0].Content; } [Theory] @@ -146,8 +146,8 @@ public void TestInterpolationMalformed() mockLogger.Setup(x => x.Error(It.IsAny())); var client = new LdAiClient(mockClient.Object); - var tracker = client.Config("foo", Context.New("key"), LdAiConfig.Disabled); - Assert.False(tracker.Config.Enabled); + var config = client.Config("foo", Context.New("key"), LdAiConfig.Disabled); + Assert.False(config.Enabled); } [Fact] diff --git a/pkgs/sdk/server-ai/test/LdAiClientTest.cs b/pkgs/sdk/server-ai/test/LdAiClientTest.cs index f28e9c9d..f204718c 100644 --- a/pkgs/sdk/server-ai/test/LdAiClientTest.cs +++ b/pkgs/sdk/server-ai/test/LdAiClientTest.cs @@ -1,4 +1,7 @@ +using System; using System.Collections.Generic; +using System.Text; +using System.Text.Json; using LaunchDarkly.Logging; using LaunchDarkly.Sdk.Server.Ai.Adapters; using LaunchDarkly.Sdk.Server.Ai.Config; @@ -16,8 +19,8 @@ public void CanInstantiateWithServerSideClient() { var client = new LdClientAdapter(new LdClient(Configuration.Builder("key").Offline(true).Build())); var aiClient = new LdAiClient(client); - var result= aiClient.CompletionConfig("foo", Context.New("key"), LdAiConfig.Disabled); - Assert.False(result.Config.Enabled); + var result = aiClient.CompletionConfig("foo", Context.New("key"), LdAiConfig.Disabled); + Assert.False(result.Enabled); } [Fact] @@ -44,9 +47,10 @@ public void ReturnsDefaultConfigWhenGivenInvalidVariation() var defaultConfig = LdAiConfig.New().AddMessage("Hello").Build(); - var tracker = client.CompletionConfig("foo", Context.New(ContextKind.Default, "key"), defaultConfig); + var config = client.CompletionConfig("foo", Context.New(ContextKind.Default, "key"), defaultConfig); - Assert.Equal(defaultConfig, tracker.Config); + Assert.Equal(defaultConfig.Enabled, config.Enabled); + Assert.Equal(defaultConfig.Messages, config.Messages); } [Fact] @@ -85,7 +89,7 @@ public void CompletionConfigMethodCallsTrackWithCorrectParameters() var client = new LdAiClient(mockClient.Object); var defaultConfig = LdAiConfig.New().Build(); - var tracker = client.CompletionConfig(configKey, context, defaultConfig); + var config = client.CompletionConfig(configKey, context, defaultConfig); mockClient.Verify(c => c.Track( "$ld:ai:usage:completion-config", @@ -93,7 +97,7 @@ public void CompletionConfigMethodCallsTrackWithCorrectParameters() LdValue.Of(configKey), 1), Times.Once); - Assert.NotNull(tracker); + Assert.NotNull(config); } [Fact] @@ -165,10 +169,10 @@ public void ConfigNotEnabledReturnsDisabledInstance(string json) // All the JSON inputs here are considered disabled, either due to lack of the 'enabled' property, // or if present, it is set to false. Therefore, if the default was returned, we'd see the assertion fail // (since calling LdAiConfig.New() constructs an enabled config by default.) - var tracker = client.CompletionConfig("foo", Context.New(ContextKind.Default, "key"), + var config = client.CompletionConfig("foo", Context.New(ContextKind.Default, "key"), LdAiConfig.New().AddMessage("foo").Build()); - Assert.False(tracker.Config.Enabled); + Assert.False(config.Enabled); } [Fact] @@ -185,7 +189,7 @@ public void CanSetAllDefaultValueFields() var client = new LdAiClient(mockClient.Object); - var tracker = client.CompletionConfig("foo", Context.New(ContextKind.Default, "key"), + var config = client.CompletionConfig("foo", Context.New(ContextKind.Default, "key"), LdAiConfig.New(). AddMessage("foo"). SetModelParam("foo", LdValue.Of("bar")). @@ -194,17 +198,17 @@ public void CanSetAllDefaultValueFields() SetModelProviderName("amazing-provider"). SetEnabled(true).Build()); - Assert.True(tracker.Config.Enabled); - Assert.Collection(tracker.Config.Messages, + Assert.True(config.Enabled); + Assert.Collection(config.Messages, message => { Assert.Equal("foo", message.Content); Assert.Equal(Role.User, message.Role); }); - Assert.Equal("amazing-provider", tracker.Config.Provider.Name); - Assert.Equal("bar", tracker.Config.Model.Parameters["foo"].AsString); - Assert.Equal("baz", tracker.Config.Model.Custom["foo"].AsString); - Assert.Equal("awesome-model", tracker.Config.Model.Name); + Assert.Equal("amazing-provider", config.Provider.Name); + Assert.Equal("bar", config.Model.Parameters["foo"].AsString); + Assert.Equal("baz", config.Model.Custom["foo"].AsString); + Assert.Equal("awesome-model", config.Model.Name); } [Fact] @@ -230,20 +234,20 @@ public void ConfigEnabledReturnsInstance() var client = new LdAiClient(mockClient.Object); // We shouldn't get this default. - var tracker = client.CompletionConfig("foo", context, + var config = client.CompletionConfig("foo", context, LdAiConfig.New().AddMessage("Goodbye!").Build()); - Assert.Collection(tracker.Config.Messages, + Assert.Collection(config.Messages, message => { Assert.Equal("Hello!", message.Content); Assert.Equal(Role.System, message.Role); }); - Assert.Equal("", tracker.Config.Provider.Name); - Assert.Equal("", tracker.Config.Model.Name); - Assert.Empty(tracker.Config.Model.Custom); - Assert.Empty(tracker.Config.Model.Parameters); + Assert.Equal("", config.Provider.Name); + Assert.Equal("", config.Model.Name); + Assert.Empty(config.Model.Custom); + Assert.Empty(config.Model.Parameters); } @@ -282,14 +286,14 @@ public void ModelParametersAreParsed() var client = new LdAiClient(mockClient.Object); // We shouldn't get this default. - var tracker = client.CompletionConfig("foo", context, + var config = client.CompletionConfig("foo", context, LdAiConfig.New().AddMessage("Goodbye!").Build()); - Assert.Equal("model-foo", tracker.Config.Model.Name); - Assert.Equal("bar", tracker.Config.Model.Parameters["foo"].AsString); - Assert.Equal(42, tracker.Config.Model.Parameters["baz"].AsInt); - Assert.Equal("baz", tracker.Config.Model.Custom["foo"].AsString); - Assert.Equal(43, tracker.Config.Model.Custom["baz"].AsInt); + Assert.Equal("model-foo", config.Model.Name); + Assert.Equal("bar", config.Model.Parameters["foo"].AsString); + Assert.Equal(42, config.Model.Parameters["baz"].AsInt); + Assert.Equal("baz", config.Model.Custom["foo"].AsString); + Assert.Equal(43, config.Model.Custom["baz"].AsInt); } [Fact] @@ -319,10 +323,10 @@ public void ProviderConfigIsParsed() var client = new LdAiClient(mockClient.Object); // We shouldn't get this default. - var tracker = client.CompletionConfig("foo", context, + var config = client.CompletionConfig("foo", context, LdAiConfig.New().AddMessage("Goodbye!").Build()); - Assert.Equal("amazing-provider", tracker.Config.Provider.Name); + Assert.Equal("amazing-provider", config.Provider.Name); } [Fact] @@ -336,9 +340,9 @@ public void ConfigWithoutDefaultValueUsesDisabledConfig() mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); var client = new LdAiClient(mockClient.Object); - var tracker = client.Config("foo", Context.New(ContextKind.Default, "key")); + var config = client.Config("foo", Context.New(ContextKind.Default, "key")); - Assert.False(tracker.Config.Enabled); + Assert.False(config.Enabled); } [Fact] @@ -355,4 +359,333 @@ public void DisabledMethodReturnsNewInstanceEachCall() var second = LdAiConfig.Disabled; Assert.NotSame(first, second); } + + [Fact] + public void CompletionConfigReturnsConfigWithCreateTrackerFactory() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + + var context = Context.New("user-key"); + const string configKey = "my-config"; + + mockClient.Setup(x => + x.JsonVariation(configKey, It.IsAny(), It.IsAny())).Returns( + LdValue.ObjectFrom(new Dictionary + { + ["_ldMeta"] = LdValue.ObjectFrom(new Dictionary + { + ["enabled"] = LdValue.Of(true), + ["variationKey"] = LdValue.Of("var-1"), + ["version"] = LdValue.Of(3) + }), + ["model"] = LdValue.ObjectFrom(new Dictionary + { + ["name"] = LdValue.Of("gpt-4") + }), + ["provider"] = LdValue.ObjectFrom(new Dictionary + { + ["name"] = LdValue.Of("openai") + }), + ["messages"] = LdValue.ArrayOf() + })); + + var client = new LdAiClient(mockClient.Object); + var config = client.CompletionConfig(configKey, context); + + var tracker = config.CreateTracker(); + Assert.NotNull(tracker); + } + + [Fact] + public void CreateTrackerFactoryReturnsTrackersWithFreshRunIds() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + + var context = Context.New("user-key"); + const string configKey = "my-config"; + + mockClient.Setup(x => + x.JsonVariation(configKey, It.IsAny(), It.IsAny())).Returns( + LdValue.ObjectFrom(new Dictionary + { + ["_ldMeta"] = LdValue.ObjectFrom(new Dictionary + { + ["enabled"] = LdValue.Of(true), + ["variationKey"] = LdValue.Of("var-1"), + ["version"] = LdValue.Of(3) + }), + ["messages"] = LdValue.ArrayOf() + })); + + var client = new LdAiClient(mockClient.Object); + var config = client.CompletionConfig(configKey, context); + + var tracker1 = config.CreateTracker(); + var tracker2 = config.CreateTracker(); + + tracker1.TrackDuration(10.0f); + tracker2.TrackDuration(20.0f); + + string runId1 = null; + string runId2 = null; + + foreach (var call in mockClient.Invocations) + { + if (call.Method.Name == "Track" && (string)call.Arguments[0] == "$ld:ai:duration:total") + { + var data = (LdValue)call.Arguments[2]; + var runId = data.Get("runId").AsString; + if (runId1 == null) runId1 = runId; + else runId2 = runId; + } + } + + Assert.NotNull(runId1); + Assert.NotNull(runId2); + Assert.NotEqual(runId1, runId2); + } + + [Fact] + public void CreateTrackerFactoryReturnsTrackersWithIndependentState() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + + var context = Context.New("user-key"); + const string configKey = "my-config"; + + mockClient.Setup(x => + x.JsonVariation(configKey, It.IsAny(), It.IsAny())).Returns( + LdValue.ObjectFrom(new Dictionary + { + ["_ldMeta"] = LdValue.ObjectFrom(new Dictionary + { + ["enabled"] = LdValue.Of(true), + ["variationKey"] = LdValue.Of("var-1"), + ["version"] = LdValue.Of(3) + }), + ["messages"] = LdValue.ArrayOf() + })); + + var client = new LdAiClient(mockClient.Object); + var config = client.CompletionConfig(configKey, context); + + var tracker1 = config.CreateTracker(); + var tracker2 = config.CreateTracker(); + + // Both should be able to track independently + tracker1.TrackSuccess(); + tracker2.TrackSuccess(); + + mockClient.Verify(x => x.Track("$ld:ai:generation:success", context, + It.IsAny(), It.IsAny()), Times.Exactly(2)); + } + + [Fact] + public void DefaultConfigGetsCreateTrackerFactory() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + + mockClient.Setup(x => + x.JsonVariation("foo", It.IsAny(), It.IsAny())).Returns(LdValue.Null); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + + var client = new LdAiClient(mockClient.Object); + var config = client.CompletionConfig("foo", Context.New("key")); + + var tracker = config.CreateTracker(); + Assert.NotNull(tracker); + } + + [Fact] + public void CreateTrackerFromResumptionTokenRoundTrips() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + + var context = Context.New("user-key"); + const string configKey = "my-config"; + + mockClient.Setup(x => + x.JsonVariation(configKey, It.IsAny(), It.IsAny())).Returns( + LdValue.ObjectFrom(new Dictionary + { + ["_ldMeta"] = LdValue.ObjectFrom(new Dictionary + { + ["enabled"] = LdValue.Of(true), + ["variationKey"] = LdValue.Of("var-1"), + ["version"] = LdValue.Of(3) + }), + ["model"] = LdValue.ObjectFrom(new Dictionary + { + ["name"] = LdValue.Of("gpt-4") + }), + ["provider"] = LdValue.ObjectFrom(new Dictionary + { + ["name"] = LdValue.Of("openai") + }), + ["messages"] = LdValue.ArrayOf() + })); + + var client = new LdAiClient(mockClient.Object); + var config = client.CompletionConfig(configKey, context); + var originalTracker = config.CreateTracker(); + var token = originalTracker.ResumptionToken; + + // Reconstruct in a different context + var newContext = Context.New("other-user"); + var resumedTracker = client.CreateTracker(token, newContext); + + Assert.NotNull(resumedTracker); + + // Track on both and verify the resumed tracker uses the same runId + originalTracker.TrackDuration(100); + resumedTracker.TrackDuration(200); + + string originalRunId = null; + string resumedRunId = null; + + foreach (var call in mockClient.Invocations) + { + if (call.Method.Name == "Track" && (string)call.Arguments[0] == "$ld:ai:duration:total") + { + var data = (LdValue)call.Arguments[2]; + var runId = data.Get("runId").AsString; + if (originalRunId == null) originalRunId = runId; + else resumedRunId = runId; + } + } + + Assert.NotNull(originalRunId); + Assert.NotNull(resumedRunId); + Assert.Equal(originalRunId, resumedRunId); + } + + [Fact] + public void CreateTrackerFromResumptionTokenSetsEmptyModelAndProvider() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + + var context = Context.New("user-key"); + + // Build a token manually with known values + var payload = JsonSerializer.Serialize(new + { + runId = "test-run-id", + configKey = "test-key", + variationKey = "var-1", + version = 2, + }); + var base64 = Convert.ToBase64String(Encoding.UTF8.GetBytes(payload)); + var token = base64.Replace('+', '-').Replace('/', '_').TrimEnd('='); + + var client = new LdAiClient(mockClient.Object); + var tracker = client.CreateTracker(token, context); + + // Track something and verify the track data has empty model/provider + tracker.TrackSuccess(); + + mockClient.Verify(x => x.Track("$ld:ai:generation:success", context, + It.Is(d => + d.Get("runId").AsString == "test-run-id" && + d.Get("configKey").AsString == "test-key" && + d.Get("variationKey").AsString == "var-1" && + d.Get("version").AsInt == 2 && + d.Get("modelName").AsString == "" && + d.Get("providerName").AsString == ""), + 1.0f), Times.Once); + } + + [Fact] + public void CreateTrackerFromInvalidTokenThrows() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + + var client = new LdAiClient(mockClient.Object); + var context = Context.New("user-key"); + + Assert.Throws(() => client.CreateTracker("not-valid-base64!!!", context)); + } + + [Fact] + public void CreateTrackerFromNullTokenThrows() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + + var client = new LdAiClient(mockClient.Object); + var context = Context.New("user-key"); + + Assert.Throws(() => client.CreateTracker(null, context)); + } + + [Fact] + public void CreateTrackerFromTokenMissingRunIdThrows() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + + var client = new LdAiClient(mockClient.Object); + var context = Context.New("user-key"); + + var payload = JsonSerializer.Serialize(new + { + configKey = "test-key", + variationKey = "var-1", + version = 1, + }); + var base64 = Convert.ToBase64String(Encoding.UTF8.GetBytes(payload)); + var token = base64.Replace('+', '-').Replace('/', '_').TrimEnd('='); + + Assert.Throws(() => client.CreateTracker(token, context)); + } + + [Fact] + public void CreateTrackerFromTokenWithoutVariationKeyHandlesAbsence() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + + var context = Context.New("user-key"); + + // Token without variationKey + var payload = JsonSerializer.Serialize(new + { + runId = "test-run-id", + configKey = "test-key", + version = 3, + }); + var base64 = Convert.ToBase64String(Encoding.UTF8.GetBytes(payload)); + var token = base64.Replace('+', '-').Replace('/', '_').TrimEnd('='); + + var client = new LdAiClient(mockClient.Object); + var tracker = client.CreateTracker(token, context); + + Assert.NotNull(tracker); + + // Track and verify variationKey is null in the track data + tracker.TrackSuccess(); + + mockClient.Verify(x => x.Track("$ld:ai:generation:success", context, + It.Is(d => + d.Get("runId").AsString == "test-run-id" && + d.Get("configKey").AsString == "test-key" && + d.Get("variationKey").IsNull && + d.Get("version").AsInt == 3), + 1.0f), Times.Once); + } } diff --git a/pkgs/sdk/server-ai/test/LdAiConfigTrackerTest.cs b/pkgs/sdk/server-ai/test/LdAiConfigTrackerTest.cs index 70c397dd..28116a9b 100644 --- a/pkgs/sdk/server-ai/test/LdAiConfigTrackerTest.cs +++ b/pkgs/sdk/server-ai/test/LdAiConfigTrackerTest.cs @@ -1,4 +1,9 @@ +using System; using System.Collections.Generic; +using System.Linq; +using System.Text; +using System.Text.Json; +using System.Threading; using System.Threading.Tasks; using LaunchDarkly.Sdk.Server.Ai.Config; using LaunchDarkly.Sdk.Server.Ai.Interfaces; @@ -10,27 +15,36 @@ namespace LaunchDarkly.Sdk.Server.Ai { public class LdAiTrackerTest { - [Fact] - public void ThrowsIfClientIsNull() + private static LdAiConfigTracker MakeTracker(Mock mockClient, string flagKey, + Context context, string variationKey = "", int version = 1, string modelName = "", + string providerName = "") { - Assert.Throws(() => - new LdAiConfigTracker(null, "key", LdAiConfig.Disabled, Context.New("key"))); + return new LdAiConfigTracker(mockClient.Object, Guid.NewGuid().ToString(), flagKey, + variationKey, version, context, modelName, providerName); } - [Fact] - public void ThrowsIfConfigIsNull() + private static bool MatchesTrackData(LdValue actual, string flagKey) { - var mockClient = new Mock(); - Assert.Throws(() => - new LdAiConfigTracker(mockClient.Object, "key", null, Context.New("key"))); + return actual.Get("configKey").Equals(LdValue.Of(flagKey)) && + actual.Get("runId").Type == LdValueType.String && + actual.Get("runId").AsString.Length > 0; } [Fact] - public void ThrowsIfKeyIsNull() + public void TrackDataIncludesRunId() { var mockClient = new Mock(); - Assert.Throws(() => - new LdAiConfigTracker(mockClient.Object, null, LdAiConfig.Disabled, Context.New("key"))); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); + + tracker.TrackDuration(1.0f); + mockClient.Verify(x => x.Track("$ld:ai:duration:total", context, + It.Is(d => + d.Get("runId").Type == LdValueType.String && + d.Get("runId").AsString.Length > 0), + 1.0f), Times.Once); } [Fact] @@ -39,19 +53,11 @@ public void CanTrackDuration() var mockClient = new Mock(); var context = Context.New("key"); const string flagKey = "key"; - var config = LdAiConfig.Disabled; - var data = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey) }, - { "version", LdValue.Of(config.Version) }, - { "configKey", LdValue.Of(flagKey) }, - { "modelName", LdValue.Of(config.Model.Name) }, - { "providerName", LdValue.Of(config.Provider.Name) } - }); - var tracker = new LdAiConfigTracker(mockClient.Object, flagKey, config, context); + var tracker = MakeTracker(mockClient, flagKey, context); tracker.TrackDuration(1.0f); - mockClient.Verify(x => x.Track("$ld:ai:duration:total", context, data, 1.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:duration:total", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); } [Fact] @@ -60,19 +66,11 @@ public void CanTrackTimeToFirstToken() var mockClient = new Mock(); var context = Context.New("key"); const string flagKey = "key"; - var config = LdAiConfig.Disabled; - var data = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey) }, - { "version", LdValue.Of(config.Version) }, - { "configKey", LdValue.Of(flagKey) }, - { "modelName", LdValue.Of(config.Model.Name) }, - { "providerName", LdValue.Of(config.Provider.Name) } - }); - var tracker = new LdAiConfigTracker(mockClient.Object, flagKey, config, context); + var tracker = MakeTracker(mockClient, flagKey, context); tracker.TrackTimeToFirstToken(1.0f); - mockClient.Verify(x => x.Track("$ld:ai:tokens:ttf", context, data, 1.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:ttf", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); } [Fact] @@ -81,19 +79,11 @@ public void CanTrackSuccess() var mockClient = new Mock(); var context = Context.New("key"); const string flagKey = "key"; - var config = LdAiConfig.Disabled; - var data = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey) }, - { "version", LdValue.Of(config.Version) }, - { "configKey", LdValue.Of(flagKey) }, - { "modelName", LdValue.Of(config.Model.Name) }, - { "providerName", LdValue.Of(config.Provider.Name) } - }); - - var tracker = new LdAiConfigTracker(mockClient.Object, flagKey, config, context); + + var tracker = MakeTracker(mockClient, flagKey, context); tracker.TrackSuccess(); - mockClient.Verify(x => x.Track("$ld:ai:generation:success", context, data, 1.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:generation:success", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); } @@ -103,19 +93,11 @@ public void CanTrackError() var mockClient = new Mock(); var context = Context.New("key"); const string flagKey = "key"; - var config = LdAiConfig.Disabled; - var data = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey) }, - { "version", LdValue.Of(config.Version) }, - { "configKey", LdValue.Of(flagKey) }, - { "modelName", LdValue.Of(config.Model.Name) }, - { "providerName", LdValue.Of(config.Provider.Name) } - }); - - var tracker = new LdAiConfigTracker(mockClient.Object, flagKey, config, context); + + var tracker = MakeTracker(mockClient, flagKey, context); tracker.TrackError(); - mockClient.Verify(x => x.Track("$ld:ai:generation:error", context, data, 1.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:generation:error", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); } @@ -125,17 +107,8 @@ public async void CanTrackDurationOfTask() var mockClient = new Mock(); var context = Context.New("key"); const string flagKey = "key"; - var config = LdAiConfig.Disabled; - var data = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey) }, - { "version", LdValue.Of(config.Version) }, - { "configKey", LdValue.Of(flagKey) }, - { "modelName", LdValue.Of(config.Model.Name) }, - { "providerName", LdValue.Of(config.Provider.Name) } - }); - var tracker = new LdAiConfigTracker(mockClient.Object, flagKey, config, context); + var tracker = MakeTracker(mockClient, flagKey, context); const int waitMs = 10; @@ -151,33 +124,38 @@ public async void CanTrackDurationOfTask() // error so this isn't flaky. If this proves to be really flaky, we can at least constrain it to be // between 0 and some large number. mockClient.Verify( - x => x.Track("$ld:ai:duration:total", context, data, - It.IsInRange(0, 500, Range.Inclusive)), Times.Once); + x => x.Track("$ld:ai:duration:total", context, + It.Is(d => MatchesTrackData(d, flagKey)), + It.IsInRange(0, 500, Moq.Range.Inclusive)), Times.Once); } [Fact] - public void CanTrackFeedback() + public void CanTrackPositiveFeedback() { var mockClient = new Mock(); var context = Context.New("key"); const string flagKey = "key"; - var config = LdAiConfig.Disabled; - var data = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey) }, - { "version", LdValue.Of(config.Version) }, - { "configKey", LdValue.Of(flagKey) }, - { "modelName", LdValue.Of(config.Model.Name) }, - { "providerName", LdValue.Of(config.Provider.Name) } - }); - - var tracker = new LdAiConfigTracker(mockClient.Object, flagKey, config, context); + + var tracker = MakeTracker(mockClient, flagKey, context); tracker.TrackFeedback(Feedback.Positive); + + mockClient.Verify(x => x.Track("$ld:ai:feedback:user:positive", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); + } + + [Fact] + public void CanTrackNegativeFeedback() + { + var mockClient = new Mock(); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); tracker.TrackFeedback(Feedback.Negative); - mockClient.Verify(x => x.Track("$ld:ai:feedback:user:positive", context, data, 1.0f), Times.Once); - mockClient.Verify(x => x.Track("$ld:ai:feedback:user:negative", context, data, 1.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:feedback:user:negative", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); } [Fact] @@ -186,17 +164,8 @@ public void CanTrackTokens() var mockClient = new Mock(); var context = Context.New("key"); const string flagKey = "key"; - var config = LdAiConfig.Disabled; - var data = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey) }, - { "version", LdValue.Of(config.Version) }, - { "configKey", LdValue.Of(flagKey) }, - { "modelName", LdValue.Of(config.Model.Name) }, - { "providerName", LdValue.Of(config.Provider.Name) } - }); - var tracker = new LdAiConfigTracker(mockClient.Object, flagKey, config, context); + var tracker = MakeTracker(mockClient, flagKey, context); var givenUsage = new Usage { @@ -206,9 +175,12 @@ public void CanTrackTokens() }; tracker.TrackTokens(givenUsage); - mockClient.Verify(x => x.Track("$ld:ai:tokens:total", context, data, 1.0f), Times.Once); - mockClient.Verify(x => x.Track("$ld:ai:tokens:input", context, data, 2.0f), Times.Once); - mockClient.Verify(x => x.Track("$ld:ai:tokens:output", context, data, 3.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:total", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:input", context, + It.Is(d => MatchesTrackData(d, flagKey)), 2.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:output", context, + It.Is(d => MatchesTrackData(d, flagKey)), 3.0f), Times.Once); } [Fact] @@ -217,17 +189,8 @@ public void CanTrackResponseWithSpecificLatency() var mockClient = new Mock(); var context = Context.New("key"); const string flagKey = "key"; - var config = LdAiConfig.Disabled; - var data = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey) }, - { "version", LdValue.Of(config.Version) }, - { "configKey", LdValue.Of(flagKey) }, - { "modelName", LdValue.Of(config.Model.Name) }, - { "providerName", LdValue.Of(config.Provider.Name) } - }); - var tracker = new LdAiConfigTracker(mockClient.Object, flagKey, config, context); + var tracker = MakeTracker(mockClient, flagKey, context); var givenUsage = new Usage { @@ -249,11 +212,16 @@ public void CanTrackResponseWithSpecificLatency() var result = tracker.TrackRequest(Task.Run(() => givenResponse)); Assert.Equal(givenResponse, result.Result); - mockClient.Verify(x => x.Track("$ld:ai:generation:success", context, data, 1.0f), Times.Once); - mockClient.Verify(x => x.Track("$ld:ai:tokens:total", context, data, 1.0f), Times.Once); - mockClient.Verify(x => x.Track("$ld:ai:tokens:input", context, data, 2.0f), Times.Once); - mockClient.Verify(x => x.Track("$ld:ai:tokens:output", context, data, 3.0f), Times.Once); - mockClient.Verify(x => x.Track("$ld:ai:duration:total", context, data, 500.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:generation:success", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:total", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:input", context, + It.Is(d => MatchesTrackData(d, flagKey)), 2.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:output", context, + It.Is(d => MatchesTrackData(d, flagKey)), 3.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:duration:total", context, + It.Is(d => MatchesTrackData(d, flagKey)), 500.0f), Times.Once); } [Fact] @@ -262,17 +230,8 @@ public void CanTrackResponseWithPartialData() var mockClient = new Mock(); var context = Context.New("key"); const string flagKey = "key"; - var config = LdAiConfig.Disabled; - var data = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey) }, - { "version", LdValue.Of(config.Version) }, - { "configKey", LdValue.Of(flagKey) }, - { "modelName", LdValue.Of(config.Model.Name) }, - { "providerName", LdValue.Of(config.Provider.Name) } - }); - var tracker = new LdAiConfigTracker(mockClient.Object, flagKey, config, context); + var tracker = MakeTracker(mockClient, flagKey, context); var givenUsage = new Usage { @@ -287,10 +246,12 @@ public void CanTrackResponseWithPartialData() var result = tracker.TrackRequest(Task.Run(() => givenResponse)); Assert.Equal(givenResponse, result.Result); - mockClient.Verify(x => x.Track("$ld:ai:tokens:total", context, data, 1.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:total", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); // if latency isn't provided via Statistics, then it is automatically measured. - mockClient.Verify(x => x.Track("$ld:ai:duration:total", context, data, It.IsAny()), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:duration:total", context, + It.Is(d => MatchesTrackData(d, flagKey)), It.IsAny()), Times.Once); } [Fact] @@ -299,24 +260,539 @@ public async Task CanTrackExceptionFromResponse() var mockClient = new Mock(); var context = Context.New("key"); const string flagKey = "key"; - var config = LdAiConfig.Disabled; - var data = LdValue.ObjectFrom(new Dictionary - { - { "variationKey", LdValue.Of(config.VariationKey) }, - { "version", LdValue.Of(config.Version) }, - { "configKey", LdValue.Of(flagKey) }, - { "modelName", LdValue.Of(config.Model.Name) }, - { "providerName", LdValue.Of(config.Provider.Name) } - }); - var tracker = new LdAiConfigTracker(mockClient.Object, flagKey, config, context); + var tracker = MakeTracker(mockClient, flagKey, context); await Assert.ThrowsAsync(() => tracker.TrackRequest(Task.FromException(new System.Exception("I am an exception")))); - mockClient.Verify(x => x.Track("$ld:ai:generation:error", context, data, 1.0f), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:generation:error", context, + It.Is(d => MatchesTrackData(d, flagKey)), 1.0f), Times.Once); // if latency isn't provided via Statistics, then it is automatically measured. - mockClient.Verify(x => x.Track("$ld:ai:duration:total", context, data, It.IsAny()), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:duration:total", context, + It.Is(d => MatchesTrackData(d, flagKey)), It.IsAny()), Times.Once); + } + + [Fact] + public void DuplicateTrackDurationIsIgnored() + { + var mockClient = new Mock(); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); + + tracker.TrackDuration(1.0f); + tracker.TrackDuration(2.0f); + + mockClient.Verify(x => x.Track("$ld:ai:duration:total", context, + It.IsAny(), It.IsAny()), Times.Once); + } + + [Fact] + public void DuplicateTrackTimeToFirstTokenIsIgnored() + { + var mockClient = new Mock(); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); + + tracker.TrackTimeToFirstToken(1.0f); + tracker.TrackTimeToFirstToken(2.0f); + + mockClient.Verify(x => x.Track("$ld:ai:tokens:ttf", context, + It.IsAny(), It.IsAny()), Times.Once); + } + + [Fact] + public void DuplicateTrackTokensIsIgnored() + { + var mockClient = new Mock(); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); + + var usage = new Usage { Total = 1, Input = 2, Output = 3 }; + + tracker.TrackTokens(usage); + tracker.TrackTokens(usage); + + mockClient.Verify(x => x.Track("$ld:ai:tokens:total", context, + It.IsAny(), It.IsAny()), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:input", context, + It.IsAny(), It.IsAny()), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:output", context, + It.IsAny(), It.IsAny()), Times.Once); + } + + [Fact] + public void TrackTokensWithEmptyUsageDoesNotConsumeAtMostOnceSlot() + { + var mockClient = new Mock(); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); + + // Empty usage emits no events. + tracker.TrackTokens(default(Usage)); + + mockClient.Verify(x => x.Track("$ld:ai:tokens:total", context, + It.IsAny(), It.IsAny()), Times.Never); + mockClient.Verify(x => x.Track("$ld:ai:tokens:input", context, + It.IsAny(), It.IsAny()), Times.Never); + mockClient.Verify(x => x.Track("$ld:ai:tokens:output", context, + It.IsAny(), It.IsAny()), Times.Never); + Assert.Null(tracker.Summary.Tokens); + + // A subsequent call with positive values should still record. + var realUsage = new Usage { Total = 100, Input = 40, Output = 60 }; + tracker.TrackTokens(realUsage); + + mockClient.Verify(x => x.Track("$ld:ai:tokens:total", context, + It.IsAny(), It.IsAny()), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:input", context, + It.IsAny(), It.IsAny()), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:tokens:output", context, + It.IsAny(), It.IsAny()), Times.Once); + + // And the slot is now consumed. + Assert.Equal(realUsage, tracker.Summary.Tokens); + } + + [Fact] + public void DuplicateTrackFeedbackIsIgnored() + { + var mockClient = new Mock(); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); + + tracker.TrackFeedback(Feedback.Positive); + tracker.TrackFeedback(Feedback.Negative); + + mockClient.Verify(x => x.Track("$ld:ai:feedback:user:positive", context, + It.IsAny(), It.IsAny()), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:feedback:user:negative", context, + It.IsAny(), It.IsAny()), Times.Never); + } + + [Fact] + public void TrackFeedbackWithInvalidValueDoesNotConsumeAtMostOnceSlot() + { + var mockClient = new Mock(); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); + + Assert.Throws(() => tracker.TrackFeedback((Feedback)42)); + + // No event emitted by the invalid call. + mockClient.Verify(x => x.Track("$ld:ai:feedback:user:positive", context, + It.IsAny(), It.IsAny()), Times.Never); + mockClient.Verify(x => x.Track("$ld:ai:feedback:user:negative", context, + It.IsAny(), It.IsAny()), Times.Never); + + // A subsequent valid call should record normally. + tracker.TrackFeedback(Feedback.Positive); + mockClient.Verify(x => x.Track("$ld:ai:feedback:user:positive", context, + It.IsAny(), It.IsAny()), Times.Once); + + // And the slot is now consumed. + Assert.Equal(Feedback.Positive, tracker.Summary.Feedback); + } + + [Fact] + public void DuplicateTrackSuccessIsIgnored() + { + var mockClient = new Mock(); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); + + tracker.TrackSuccess(); + tracker.TrackSuccess(); + + mockClient.Verify(x => x.Track("$ld:ai:generation:success", context, + It.IsAny(), It.IsAny()), Times.Once); + } + + [Fact] + public void TrackErrorAfterSuccessIsIgnored() + { + var mockClient = new Mock(); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); + + tracker.TrackSuccess(); + tracker.TrackError(); + + mockClient.Verify(x => x.Track("$ld:ai:generation:success", context, + It.IsAny(), It.IsAny()), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:generation:error", context, + It.IsAny(), It.IsAny()), Times.Never); + } + + [Fact] + public void TrackSuccessAfterErrorIsIgnored() + { + var mockClient = new Mock(); + var context = Context.New("key"); + const string flagKey = "key"; + + var tracker = MakeTracker(mockClient, flagKey, context); + + tracker.TrackError(); + tracker.TrackSuccess(); + + mockClient.Verify(x => x.Track("$ld:ai:generation:error", context, + It.IsAny(), It.IsAny()), Times.Once); + mockClient.Verify(x => x.Track("$ld:ai:generation:success", context, + It.IsAny(), It.IsAny()), Times.Never); + } + + [Fact] + public void ResumptionTokenContainsExpectedFields() + { + var mockClient = new Mock(); + var context = Context.New("key"); + + var tracker = new LdAiConfigTracker(mockClient.Object, Guid.NewGuid().ToString(), + "my-config-key", "", 1, context, "test-model", "test-provider"); + var token = tracker.ResumptionToken; + + Assert.NotNull(token); + Assert.NotEmpty(token); + + // Decode and verify the payload + var base64 = token.Replace('-', '+').Replace('_', '/'); + switch (base64.Length % 4) + { + case 2: base64 += "=="; break; + case 3: base64 += "="; break; + } + var json = Encoding.UTF8.GetString(Convert.FromBase64String(base64)); + var doc = JsonDocument.Parse(json); + + Assert.Equal("my-config-key", doc.RootElement.GetProperty("configKey").GetString()); + Assert.Equal(1, doc.RootElement.GetProperty("version").GetInt32()); + Assert.True(doc.RootElement.GetProperty("runId").GetString().Length > 0); + + // variationKey is empty, so it should be omitted + Assert.False(doc.RootElement.TryGetProperty("variationKey", out _)); + + // modelName and providerName should NOT be in the token + Assert.False(doc.RootElement.TryGetProperty("modelName", out _)); + Assert.False(doc.RootElement.TryGetProperty("providerName", out _)); + } + + [Fact] + public void ResumptionTokenHasCanonicalKeyOrder() + { + var mockClient = new Mock(); + var context = Context.New("key"); + + var tracker = new LdAiConfigTracker(mockClient.Object, "test-run-id", + "my-config-key", "my-variation", 5, context, "", ""); + var token = tracker.ResumptionToken; + + var base64 = token.Replace('-', '+').Replace('_', '/'); + switch (base64.Length % 4) + { + case 2: base64 += "=="; break; + case 3: base64 += "="; break; + } + var json = Encoding.UTF8.GetString(Convert.FromBase64String(base64)); + + // Verify canonical order: runId → configKey → variationKey → version + var runIdIdx = json.IndexOf("\"runId\"", StringComparison.Ordinal); + var configKeyIdx = json.IndexOf("\"configKey\"", StringComparison.Ordinal); + var variationKeyIdx = json.IndexOf("\"variationKey\"", StringComparison.Ordinal); + var versionIdx = json.IndexOf("\"version\"", StringComparison.Ordinal); + + Assert.True(runIdIdx < configKeyIdx); + Assert.True(configKeyIdx < variationKeyIdx); + Assert.True(variationKeyIdx < versionIdx); + } + + [Fact] + public void ResumptionTokenIsUrlSafeBase64() + { + var mockClient = new Mock(); + var context = Context.New("key"); + + var tracker = new LdAiConfigTracker(mockClient.Object, Guid.NewGuid().ToString(), + "key", "", 1, context, "", ""); + var token = tracker.ResumptionToken; + + // URL-safe base64 should not contain +, /, or = + Assert.DoesNotContain("+", token); + Assert.DoesNotContain("/", token); + Assert.DoesNotContain("=", token); + } + + [Fact] + public void ResumptionTokenIsConsistentAcrossCalls() + { + var mockClient = new Mock(); + var context = Context.New("key"); + + var tracker = new LdAiConfigTracker(mockClient.Object, Guid.NewGuid().ToString(), + "key", "", 1, context, "", ""); + + var token1 = tracker.ResumptionToken; + var token2 = tracker.ResumptionToken; + + Assert.Equal(token1, token2); + } + + [Fact] + public void ResumptionTokenIncludesVariationKeyWhenPresent() + { + var mockClient = new Mock(); + var mockLogger = new Mock(); + mockClient.Setup(x => x.GetLogger()).Returns(mockLogger.Object); + var context = Context.New("key"); + + // Use the LdAiClient to get a config with a real variationKey from flag evaluation + mockClient.Setup(x => + x.JsonVariation("key", It.IsAny(), It.IsAny())).Returns( + LdValue.ObjectFrom(new Dictionary + { + ["_ldMeta"] = LdValue.ObjectFrom(new Dictionary + { + ["enabled"] = LdValue.Of(true), + ["variationKey"] = LdValue.Of("my-variation"), + ["version"] = LdValue.Of(5) + }), + ["messages"] = LdValue.ArrayOf() + })); + + var client = new LdAiClient(mockClient.Object); + var config = client.CompletionConfig("key", context); + var tracker = config.CreateTracker(); + var token = tracker.ResumptionToken; + + // Decode and verify variationKey is present + var base64 = token.Replace('-', '+').Replace('_', '/'); + switch (base64.Length % 4) + { + case 2: base64 += "=="; break; + case 3: base64 += "="; break; + } + var json = Encoding.UTF8.GetString(Convert.FromBase64String(base64)); + var doc = JsonDocument.Parse(json); + + Assert.Equal("my-variation", doc.RootElement.GetProperty("variationKey").GetString()); + } + + [Fact] + public void FromResumptionTokenReconstructsTrackerWithOriginalRunId() + { + var mockClient = new Mock(); + var context = Context.New("key"); + + var original = new LdAiConfigTracker(mockClient.Object, Guid.NewGuid().ToString(), + "my-key", "", 1, context, "", ""); + var token = original.ResumptionToken; + + var newContext = Context.New("other-key"); + var resumed = LdAiConfigTracker.FromResumptionToken(token, mockClient.Object, newContext); + + // Both should track with the same runId + original.TrackDuration(10); + resumed.TrackDuration(20); + + string originalRunId = null; + string resumedRunId = null; + foreach (var call in mockClient.Invocations) + { + if (call.Method.Name == "Track" && (string)call.Arguments[0] == "$ld:ai:duration:total") + { + var data = (LdValue)call.Arguments[2]; + if (originalRunId == null) originalRunId = data.Get("runId").AsString; + else resumedRunId = data.Get("runId").AsString; + } + } + + Assert.NotNull(originalRunId); + Assert.Equal(originalRunId, resumedRunId); + } + + [Fact] + public void FromResumptionTokenThrowsOnNullToken() + { + var mockClient = new Mock(); + Assert.Throws(() => + LdAiConfigTracker.FromResumptionToken(null, mockClient.Object, Context.New("key"))); + } + + [Fact] + public void FromResumptionTokenThrowsOnInvalidToken() + { + var mockClient = new Mock(); + Assert.Throws(() => + LdAiConfigTracker.FromResumptionToken("not-valid!!!", mockClient.Object, Context.New("key"))); + } + + [Fact] + public void FromResumptionTokenRejectsNullJsonPayload() + { + var mockClient = new Mock(); + + // Base64-encode the JSON literal "null" so deserialization succeeds and yields a null payload. + var base64 = Convert.ToBase64String(Encoding.UTF8.GetBytes("null")); + var token = base64.Replace('+', '-').Replace('/', '_').TrimEnd('='); + + Assert.Throws(() => + LdAiConfigTracker.FromResumptionToken(token, mockClient.Object, Context.New("key"))); + } + + [Fact] + public void SummaryReflectsTrackedMetrics() + { + var mockClient = new Mock(); + var context = Context.New("key"); + + var tracker = MakeTracker(mockClient, "key", context); + + // Initially all null + var summary = tracker.Summary; + Assert.Null(summary.DurationMs); + Assert.Null(summary.Feedback); + Assert.Null(summary.Tokens); + Assert.Null(summary.Success); + Assert.Null(summary.TimeToFirstTokenMs); + + // Track some metrics + tracker.TrackDuration(100.5f); + tracker.TrackTimeToFirstToken(10.2f); + tracker.TrackFeedback(Feedback.Positive); + tracker.TrackSuccess(); + tracker.TrackTokens(new Usage { Total = 10, Input = 3, Output = 7 }); + + summary = tracker.Summary; + Assert.Equal(100.5, summary.DurationMs.Value, 1); + Assert.Equal(Feedback.Positive, summary.Feedback); + Assert.NotNull(summary.Tokens); + Assert.Equal(10, summary.Tokens.Value.Total); + Assert.Equal(3, summary.Tokens.Value.Input); + Assert.Equal(7, summary.Tokens.Value.Output); + Assert.True(summary.Success); + Assert.Equal(10.2, summary.TimeToFirstTokenMs.Value, 1); + } + + [Fact] + public void SummaryReflectsErrorState() + { + var mockClient = new Mock(); + var context = Context.New("key"); + + var tracker = MakeTracker(mockClient, "key", context); + tracker.TrackError(); + + var summary = tracker.Summary; + Assert.False(summary.Success); + } + + [Fact] + public void TrackSuccessIsAtMostOnceUnderConcurrency() + { + const int iterations = 500; + var maxObservedInvocations = 0; + var iterationWithMax = -1; + + for (var i = 0; i < iterations; i++) + { + var mockClient = new Mock(); + var context = Context.New("key"); + var tracker = MakeTracker(mockClient, "key", context); + + var barrier = new Barrier(2); + var t1 = Task.Run(() => { barrier.SignalAndWait(); tracker.TrackSuccess(); }); + var t2 = Task.Run(() => { barrier.SignalAndWait(); tracker.TrackSuccess(); }); + Task.WaitAll(t1, t2); + + var calls = mockClient.Invocations.Count(inv => + inv.Method.Name == "Track" && + (string)inv.Arguments[0] == "$ld:ai:generation:success"); + + if (calls > maxObservedInvocations) { maxObservedInvocations = calls; iterationWithMax = i; } + } + + Assert.True(maxObservedInvocations == 1, + $"Expected at-most-once emission of $ld:ai:generation:success across {iterations} " + + $"iterations, but observed {maxObservedInvocations} emissions on iteration {iterationWithMax}. " + + "TrackSuccess's check-then-set on _trackedSuccess is not atomic."); + } + + [Fact] + public void TrackSuccessAndTrackErrorAreMutuallyExclusiveUnderConcurrency() + { + const int iterations = 500; + var iterationsWithBoth = 0; + + for (var i = 0; i < iterations; i++) + { + var mockClient = new Mock(); + var context = Context.New("key"); + var tracker = MakeTracker(mockClient, "key", context); + + var barrier = new Barrier(2); + var t1 = Task.Run(() => { barrier.SignalAndWait(); tracker.TrackSuccess(); }); + var t2 = Task.Run(() => { barrier.SignalAndWait(); tracker.TrackError(); }); + Task.WaitAll(t1, t2); + + var successCalls = mockClient.Invocations.Count(inv => + inv.Method.Name == "Track" && (string)inv.Arguments[0] == "$ld:ai:generation:success"); + var errorCalls = mockClient.Invocations.Count(inv => + inv.Method.Name == "Track" && (string)inv.Arguments[0] == "$ld:ai:generation:error"); + + if (successCalls > 0 && errorCalls > 0) iterationsWithBoth++; + } + + Assert.True(iterationsWithBoth == 0, + $"Expected TrackSuccess and TrackError to be mutually exclusive across {iterations} " + + $"iterations, but {iterationsWithBoth} iterations emitted both success AND error " + + "events for the same runId. The shared _trackedSuccess slot is not atomically guarded."); + } + + [Fact] + public void TrackDurationIsAtMostOnceUnderConcurrency() + { + const int iterations = 500; + var maxObservedInvocations = 0; + var iterationWithMax = -1; + + for (var i = 0; i < iterations; i++) + { + var mockClient = new Mock(); + var context = Context.New("key"); + var tracker = MakeTracker(mockClient, "key", context); + + var barrier = new Barrier(2); + var t1 = Task.Run(() => { barrier.SignalAndWait(); tracker.TrackDuration(123.0f); }); + var t2 = Task.Run(() => { barrier.SignalAndWait(); tracker.TrackDuration(123.0f); }); + Task.WaitAll(t1, t2); + + var calls = mockClient.Invocations.Count(inv => + inv.Method.Name == "Track" && + (string)inv.Arguments[0] == "$ld:ai:duration:total"); + + if (calls > maxObservedInvocations) { maxObservedInvocations = calls; iterationWithMax = i; } + } + + Assert.True(maxObservedInvocations == 1, + $"Expected at-most-once emission of $ld:ai:duration:total across {iterations} " + + $"iterations, but observed {maxObservedInvocations} emissions on iteration {iterationWithMax}. " + + "TrackDuration's check-then-set on _durationMs is not atomic."); } } }