AI Configuration request tracing

package.json

@@ -25,7 +25,7 @@
     "dev": "rollup --config --watch",
     "lint": "eslint src/ test/",
     "fix-lint": "eslint src/ test/ --fix",
-    "test": "mocha out/test/*.test.{js,cjs,mjs} --parallel"
+    "test": "mocha out/test/requestTracing.test.{js,cjs,mjs} --parallel"
   },
   "repository": {
     "type": "git",

src/AzureAppConfigurationImpl.ts

@@ -24,11 +24,13 @@ import {
     CONDITIONS_KEY_NAME,
     CLIENT_FILTERS_KEY_NAME
 } from "./featureManagement/constants.js";
-import { FM_PACKAGE_NAME } from "./requestTracing/constants.js";
+import { FM_PACKAGE_NAME, AI_MIME_PROFILE, AI_CHAT_COMPLETION_MIME_PROFILE } from "./requestTracing/constants.js";
+import { parseContentType, isJsonContentType } from "./common/contentType.js";
 import { AzureKeyVaultKeyValueAdapter } from "./keyvault/AzureKeyVaultKeyValueAdapter.js";
 import { RefreshTimer } from "./refresh/RefreshTimer.js";
 import { RequestTracingOptions, getConfigurationSettingWithTrace, listConfigurationSettingsWithTrace, requestTracingEnabled } from "./requestTracing/utils.js";
 import { FeatureFlagTracingOptions } from "./requestTracing/FeatureFlagTracingOptions.js";
+import { AIConfigurationTracingOptions } from "./requestTracing/AIConfigurationTracingOptions.js";
 import { KeyFilter, LabelFilter, SettingSelector } from "./types.js";
 import { ConfigurationClientManager } from "./ConfigurationClientManager.js";
 
@@ -58,6 +60,7 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
     #isFailoverRequest: boolean = false;
     #featureFlagTracing: FeatureFlagTracingOptions | undefined;
     #fmVersion: string | undefined;
+    #aiConfigurationTracing: AIConfigurationTracingOptions | undefined;
 
     // Refresh
     #refreshInProgress: boolean = false;
@@ -97,6 +100,7 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
         // enable request tracing if not opt-out
         this.#requestTracingEnabled = requestTracingEnabled();
         if (this.#requestTracingEnabled) {
+            this.#aiConfigurationTracing = new AIConfigurationTracingOptions();
             this.#featureFlagTracing = new FeatureFlagTracingOptions();
         }
 
@@ -178,7 +182,8 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
             replicaCount: this.#clientManager.getReplicaCount(),
             isFailoverRequest: this.#isFailoverRequest,
             featureFlagTracing: this.#featureFlagTracing,
-            fmVersion: this.#fmVersion
+            fmVersion: this.#fmVersion,
+            aiConfigurationTracing: this.#aiConfigurationTracing
         };
     }
 
@@ -416,9 +421,14 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
             await this.#updateWatchedKeyValuesEtag(loadedSettings);
         }
 
+        if (this.#requestTracingEnabled && this.#aiConfigurationTracing !== undefined) {
+            // Reset old AI configuration tracing in order to track the information present in the current response from server.
+            this.#aiConfigurationTracing.reset();
+        }
+
         // process key-values, watched settings have higher priority
         for (const setting of loadedSettings) {
-            const [key, value] = await this.#processKeyValues(setting);
+            const [key, value] = await this.#processKeyValue(setting);
             keyValues.push([key, value]);
         }
 
@@ -467,6 +477,11 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
         const loadFeatureFlag = true;
         const featureFlagSettings = await this.#loadConfigurationSettings(loadFeatureFlag);
 
+        if (this.#requestTracingEnabled && this.#featureFlagTracing !== undefined) {
+            // Reset old feature flag tracing in order to track the information present in the current response from server.
+            this.#featureFlagTracing.reset();
+        }
+
         // parse feature flags
         const featureFlags = await Promise.all(
             featureFlagSettings.map(setting => this.#parseFeatureFlag(setting))
@@ -633,12 +648,32 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
         throw new Error("All clients failed to get configuration settings.");
     }
 
-    async #processKeyValues(setting: ConfigurationSetting<string>): Promise<[string, unknown]> {
+    async #processKeyValue(setting: ConfigurationSetting<string>): Promise<[string, unknown]> {
+        this.#setAIConfigurationTracing(setting);
+
         const [key, value] = await this.#processAdapters(setting);
         const trimmedKey = this.#keyWithPrefixesTrimmed(key);
         return [trimmedKey, value];
     }
 
+    #setAIConfigurationTracing(setting: ConfigurationSetting<string>): void {
+        if (this.#requestTracingEnabled && this.#aiConfigurationTracing !== undefined) {
+            const contentType = parseContentType(setting.contentType);
+            if (isJsonContentType(contentType)) {
+                const profile = contentType?.parameters["profile"];
+                if (profile === undefined) {
+                    return;
+                }
+                if (profile.includes(AI_MIME_PROFILE)) {
+                    this.#aiConfigurationTracing.usesAIConfiguration = true;
+                }
+                if (profile.includes(AI_CHAT_COMPLETION_MIME_PROFILE)) {
+                    this.#aiConfigurationTracing.usesAIChatCompletionConfiguration = true;
+                }
+            }
+        }
+    }
+
     async #processAdapters(setting: ConfigurationSetting<string>): Promise<[string, unknown]> {
         for (const adapter of this.#adapters) {
             if (adapter.canProcess(setting)) {
@@ -675,6 +710,20 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
             };
         }
 
+        this.#setFeatureFlagTracing(featureFlag);
+
+        return featureFlag;
+    }
+
+    #createFeatureFlagReference(setting: ConfigurationSetting<string>): string {
+        let featureFlagReference = `${this.#clientManager.endpoint.origin}/kv/${setting.key}`;
+        if (setting.label && setting.label.trim().length !== 0) {
+            featureFlagReference += `?label=${setting.label}`;
+        }
+        return featureFlagReference;
+    }
+
+    #setFeatureFlagTracing(featureFlag: any): void {
         if (this.#requestTracingEnabled && this.#featureFlagTracing !== undefined) {
             if (featureFlag[CONDITIONS_KEY_NAME] &&
                 featureFlag[CONDITIONS_KEY_NAME][CLIENT_FILTERS_KEY_NAME] &&
@@ -693,16 +742,6 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
                 this.#featureFlagTracing.usesSeed = true;
             }
         }
-
-        return featureFlag;
-    }
-
-    #createFeatureFlagReference(setting: ConfigurationSetting<string>): string {
-        let featureFlagReference = `${this.#clientManager.endpoint.origin}/kv/${setting.key}`;
-        if (setting.label && setting.label.trim().length !== 0) {
-            featureFlagReference += `?label=${setting.label}`;
-        }
-        return featureFlagReference;
     }
 }
 

src/JsonKeyValueAdapter.ts

@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 import { ConfigurationSetting, featureFlagContentType, secretReferenceContentType } from "@azure/app-configuration";
+import { parseContentType, isJsonContentType } from "./common/contentType.js";
 import { IKeyValueAdapter } from "./IKeyValueAdapter.js";
 
 export class JsonKeyValueAdapter implements IKeyValueAdapter {
@@ -17,7 +18,8 @@ export class JsonKeyValueAdapter implements IKeyValueAdapter {
         if (JsonKeyValueAdapter.#ExcludedJsonContentTypes.includes(setting.contentType)) {
             return false;
         }
-        return isJsonContentType(setting.contentType);
+        const contentType = parseContentType(setting.contentType);
+        return isJsonContentType(contentType);
     }
 
     async processKeyValue(setting: ConfigurationSetting): Promise<[string, unknown]> {
@@ -34,24 +36,3 @@ export class JsonKeyValueAdapter implements IKeyValueAdapter {
         return [setting.key, parsedValue];
     }
 }
-
-// Determine whether a content type string is a valid JSON content type.
-// https://docs.microsoft.com/en-us/azure/azure-app-configuration/howto-leverage-json-content-type
-function isJsonContentType(contentTypeValue: string): boolean {
-    if (!contentTypeValue) {
-        return false;
-    }
-
-    const contentTypeNormalized: string = contentTypeValue.trim().toLowerCase();
-    const mimeType: string = contentTypeNormalized.split(";", 1)[0].trim();
-    const typeParts: string[] = mimeType.split("/");
-    if (typeParts.length !== 2) {
-        return false;
-    }
-
-    if (typeParts[0] !== "application") {
-        return false;
-    }
-
-    return typeParts[1].split("+").includes("json");
-}

src/common/contentType.ts

@@ -0,0 +1,44 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+export type ContentType = {
+    mediaType: string;
+    parameters: Record<string, string>;
+}
+
+export function parseContentType(contentTypeValue: string | undefined): ContentType | undefined {
+    if (!contentTypeValue) {
+        return undefined;
+    }
+    const [mediaType, ...args] = contentTypeValue.split(";").map((s) => s.trim());
+    const parameters: Record<string, string> = {};
+
+    for (const param of args) {
+        const [key, value] = param.split("=").map((s) => s.trim());
+        if (key && value) {
+            parameters[key] = value;
+        }
+    }
+
+    return { mediaType, parameters };
+}
+
+// Determine whether a content type string is a valid JSON content type.
+// https://docs.microsoft.com/en-us/azure/azure-app-configuration/howto-leverage-json-content-type
+export function isJsonContentType(contentType: ContentType | undefined): boolean {
+    const mediaType = contentType?.mediaType?.trim().toLowerCase();
+    if (!mediaType) {
+        return false;
+    }
+
+    const typeParts: string[] = mediaType.split("/");
+    if (typeParts.length !== 2) {
+        return false;
+    }
+
+    if (typeParts[0] !== "application") {
+        return false;
+    }
+
+    return typeParts[1].split("+").includes("json");
+}

src/requestTracing/AIConfigurationTracingOptions.ts

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+export class AIConfigurationTracingOptions {
+    usesAIConfiguration: boolean = false;
+    usesAIChatCompletionConfiguration: boolean = false;
+
+    reset(): void {
+        this.usesAIConfiguration = false;
+        this.usesAIChatCompletionConfiguration = false;
+    }
+
+    usesAnyTracingFeature() {
+        return this.usesAIConfiguration || this.usesAIChatCompletionConfiguration;
+    }
+}

src/requestTracing/FeatureFlagTracingOptions.ts

@@ -18,7 +18,7 @@ export class FeatureFlagTracingOptions {
     usesSeed: boolean = false;
     maxVariants: number = 0;
 
-    resetFeatureFlagTracing(): void {
+    reset(): void {
         this.usesCustomFilter = false;
         this.usesTimeWindowFilter = false;
         this.usesTargetingFilter = false;

src/requestTracing/constants.ts

@@ -70,4 +70,11 @@ export const FF_MAX_VARIANTS_KEY = "MaxVariants";
 export const FF_SEED_USED_TAG = "Seed";
 export const FF_FEATURES_KEY = "FFFeatures";
 
+// AI Configuration tracing
+export const AI_CONFIGURATION_TAG = "AI";
+export const AI_CHAT_COMPLETION_CONFIGURATION_TAG = "AICC";
+
+export const AI_MIME_PROFILE = "https://azconfig.io/mime-profiles/ai";
+export const AI_CHAT_COMPLETION_MIME_PROFILE = "https://azconfig.io/mime-profiles/ai/chat-completion";
+
 export const DELIMITER = "+";

src/requestTracing/utils.ts

@@ -4,6 +4,7 @@
 import { AppConfigurationClient, ConfigurationSettingId, GetConfigurationSettingOptions, ListConfigurationSettingsOptions } from "@azure/app-configuration";
 import { AzureAppConfigurationOptions } from "../AzureAppConfigurationOptions.js";
 import { FeatureFlagTracingOptions } from "./FeatureFlagTracingOptions.js";
+import { AIConfigurationTracingOptions } from "./AIConfigurationTracingOptions.js";
 import {
     AZURE_FUNCTION_ENV_VAR,
     AZURE_WEB_APP_ENV_VAR,
@@ -28,7 +29,10 @@ import {
     FAILOVER_REQUEST_TAG,
     FEATURES_KEY,
     LOAD_BALANCE_CONFIGURED_TAG,
-    FM_VERSION_KEY
+    FM_VERSION_KEY,
+    DELIMITER,
+    AI_CONFIGURATION_TAG,
+    AI_CHAT_COMPLETION_CONFIGURATION_TAG
 } from "./constants";
 
 export interface RequestTracingOptions {
@@ -39,6 +43,7 @@ export interface RequestTracingOptions {
     isFailoverRequest: boolean;
     featureFlagTracing: FeatureFlagTracingOptions | undefined;
     fmVersion: string | undefined;
+    aiConfigurationTracing: AIConfigurationTracingOptions | undefined;
 }
 
 // Utils
@@ -125,9 +130,22 @@ export function createCorrelationContextHeader(requestTracingOptions: RequestTra
         keyValues.set(FM_VERSION_KEY, requestTracingOptions.fmVersion);
     }
 
-    // Compact tags: Features=LB+...
-    if (appConfigOptions?.loadBalancingEnabled) {
-        keyValues.set(FEATURES_KEY, LOAD_BALANCE_CONFIGURED_TAG);
+    // Compact tags: Features=LB+AI+AICC...
+    if (appConfigOptions?.loadBalancingEnabled || requestTracingOptions.aiConfigurationTracing?.usesAnyTracingFeature()) {
+        const tags: string[] = [];
+        if (appConfigOptions?.loadBalancingEnabled) {
+            tags.push(LOAD_BALANCE_CONFIGURED_TAG);
+        }
+        if (requestTracingOptions.aiConfigurationTracing?.usesAIConfiguration) {
+            tags.push(AI_CONFIGURATION_TAG);
+        }
+        if (requestTracingOptions.aiConfigurationTracing?.usesAIChatCompletionConfiguration) {
+            tags.push(AI_CHAT_COMPLETION_CONFIGURATION_TAG);
+        }
+
+        if (tags.length > 0) {
+            keyValues.set(FEATURES_KEY, tags.join(DELIMITER));
+        }
     }
 
     const contextParts: string[] = [];

test/requestTracing.test.ts

@@ -64,6 +64,19 @@ describe("request tracing", function () {
         expect(correlationContext.includes("UsesKeyVault")).eq(true);
     });
 
+    it("should have loadbalancing tag in correlation-context header", async () => {
+        try {
+            await load(createMockedConnectionString(fakeEndpoint), {
+                clientOptions,
+                loadBalancingEnabled: true,
+            });
+        } catch (e) { /* empty */ }
+        expect(headerPolicy.headers).not.undefined;
+        const correlationContext = headerPolicy.headers.get("Correlation-Context");
+        expect(correlationContext).not.undefined;
+        expect(correlationContext.includes("Features=LB")).eq(true);
+    });
+
     it("should have replica count in correlation-context header", async () => {
         const replicaCount = 2;
         sinon.stub(ConfigurationClientManager.prototype, "getReplicaCount").returns(replicaCount);
@@ -293,6 +306,37 @@ describe("request tracing", function () {
         restoreMocks();
     });
 
+    it("should have AI tag in correlation-context header if key values use AI configuration", async () => {
+        let correlationContext: string = "";
+        const listKvCallback = (listOptions) => {
+            correlationContext = listOptions?.requestOptions?.customHeaders[CORRELATION_CONTEXT_HEADER_NAME] ?? "";
+        };
+
+        mockAppConfigurationClientListConfigurationSettings([[
+            createMockedKeyValue({ contentType: "application/json; profile=\"https://azconfig.io/mime-profiles/ai/chat-completion\"" })
+        ]], listKvCallback);
+
+        const settings = await load(createMockedConnectionString(fakeEndpoint), {
+            refreshOptions: {
+                enabled: true,
+                refreshIntervalInMs: 1000
+            }
+        });
+
+        expect(correlationContext).not.undefined;
+        expect(correlationContext?.includes("RequestType=Startup")).eq(true);
+
+        await sleepInMs(1000 + 1);
+        try {
+            await settings.refresh();
+        } catch (e) { /* empty */ }
+        expect(headerPolicy.headers).not.undefined;
+        expect(correlationContext).not.undefined;
+        expect(correlationContext?.includes("Features=AI+AICC")).eq(true);
+
+        restoreMocks();
+    });
+
     describe("request tracing in Web Worker environment", () => {
         let originalNavigator;
         let originalWorkerNavigator;