Skip to content

Commit 76579dd

Browse files
zhiyuanliang-mszhiyuanliang-ms
committed
add request tracing for ai configuration usage
1 parent a5b5bb9 commit 76579dd

File tree

7 files changed

+143
-36
lines changed

7 files changed

+143
-36
lines changed

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
"dev": "rollup --config --watch",
2626
"lint": "eslint src/ test/",
2727
"fix-lint": "eslint src/ test/ --fix",
28-
"test": "mocha out/test/*.test.{js,cjs,mjs} --parallel"
28+
"test": "mocha out/test/requestTracing.test.{js,cjs,mjs} --parallel"
2929
},
3030
"repository": {
3131
"type": "git",

‎src/AzureAppConfigurationImpl.ts‎

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import {
2525
CLIENT_FILTERS_KEY_NAME
2626
} from "./featureManagement/constants.js";
2727
import { FM_PACKAGE_NAME, AI_MIME_PROFILE, AI_CHAT_COMPLETION_MIME_PROFILE } from "./requestTracing/constants.js";
28+
import { parseContentType, isJsonContentType } from "./common/contentType.js";
2829
import { AzureKeyVaultKeyValueAdapter } from "./keyvault/AzureKeyVaultKeyValueAdapter.js";
2930
import { RefreshTimer } from "./refresh/RefreshTimer.js";
3031
import { RequestTracingOptions, getConfigurationSettingWithTrace, listConfigurationSettingsWithTrace, requestTracingEnabled } from "./requestTracing/utils.js";
@@ -99,6 +100,7 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
99100
// enable request tracing if not opt-out
100101
this.#requestTracingEnabled = requestTracingEnabled();
101102
if (this.#requestTracingEnabled) {
103+
this.#aiConfigurationTracing = new AIConfigurationTracingOptions();
102104
this.#featureFlagTracing = new FeatureFlagTracingOptions();
103105
}
104106

@@ -419,9 +421,14 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
419421
await this.#updateWatchedKeyValuesEtag(loadedSettings);
420422
}
421423

424+
if (this.#requestTracingEnabled && this.#aiConfigurationTracing !== undefined) {
425+
// Reset old AI configuration tracing in order to track the information present in the current response from server.
426+
this.#aiConfigurationTracing.reset();
427+
}
428+
422429
// process key-values, watched settings have higher priority
423430
for (const setting of loadedSettings) {
424-
const [key, value] = await this.#processKeyValues(setting);
431+
const [key, value] = await this.#processKeyValue(setting);
425432
keyValues.push([key, value]);
426433
}
427434

@@ -470,6 +477,11 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
470477
const loadFeatureFlag = true;
471478
const featureFlagSettings = await this.#loadConfigurationSettings(loadFeatureFlag);
472479

480+
if (this.#requestTracingEnabled && this.#featureFlagTracing !== undefined) {
481+
// Reset old feature flag tracing in order to track the information present in the current response from server.
482+
this.#featureFlagTracing.reset();
483+
}
484+
473485
// parse feature flags
474486
const featureFlags = await Promise.all(
475487
featureFlagSettings.map(setting => this.#parseFeatureFlag(setting))
@@ -636,7 +648,7 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
636648
throw new Error("All clients failed to get configuration settings.");
637649
}
638650

639-
async #processKeyValues(setting: ConfigurationSetting<string>): Promise<[string, unknown]> {
651+
async #processKeyValue(setting: ConfigurationSetting<string>): Promise<[string, unknown]> {
640652
this.#setAIConfigurationTracing(setting);
641653

642654
const [key, value] = await this.#processAdapters(setting);
@@ -646,10 +658,19 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
646658

647659
#setAIConfigurationTracing(setting: ConfigurationSetting<string>): void {
648660
if (this.#requestTracingEnabled && this.#aiConfigurationTracing !== undefined) {
649-
// Reset old AI configuration tracing in order to track the information present in the current response from server.
650-
this.#aiConfigurationTracing.reset();
651-
652-
661+
const contentType = parseContentType(setting.contentType);
662+
if (isJsonContentType(contentType)) {
663+
const profile = contentType?.parameters["profile"];
664+
if (profile === undefined) {
665+
return;
666+
}
667+
if (profile.includes(AI_MIME_PROFILE)) {
668+
this.#aiConfigurationTracing.usesAIConfiguration = true;
669+
}
670+
if (profile.includes(AI_CHAT_COMPLETION_MIME_PROFILE)) {
671+
this.#aiConfigurationTracing.usesAIChatCompletionConfiguration = true;
672+
}
673+
}
653674
}
654675
}
655676

@@ -704,9 +725,6 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
704725

705726
#setFeatureFlagTracing(featureFlag: any): void {
706727
if (this.#requestTracingEnabled && this.#featureFlagTracing !== undefined) {
707-
// Reset old feature flag tracing in order to track the information present in the current response from server.
708-
this.#featureFlagTracing.reset();
709-
710728
if (featureFlag[CONDITIONS_KEY_NAME] &&
711729
featureFlag[CONDITIONS_KEY_NAME][CLIENT_FILTERS_KEY_NAME] &&
712730
Array.isArray(featureFlag[CONDITIONS_KEY_NAME][CLIENT_FILTERS_KEY_NAME])) {

‎src/JsonKeyValueAdapter.ts‎

Lines changed: 3 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// Licensed under the MIT license.
33

44
import { ConfigurationSetting, featureFlagContentType, secretReferenceContentType } from "@azure/app-configuration";
5+
import { parseContentType, isJsonContentType } from "./common/contentType.js";
56
import { IKeyValueAdapter } from "./IKeyValueAdapter.js";
67

78
export class JsonKeyValueAdapter implements IKeyValueAdapter {
@@ -17,7 +18,8 @@ export class JsonKeyValueAdapter implements IKeyValueAdapter {
1718
if (JsonKeyValueAdapter.#ExcludedJsonContentTypes.includes(setting.contentType)) {
1819
return false;
1920
}
20-
return isJsonContentType(setting.contentType);
21+
const contentType = parseContentType(setting.contentType);
22+
return isJsonContentType(contentType);
2123
}
2224

2325
async processKeyValue(setting: ConfigurationSetting): Promise<[string, unknown]> {
@@ -34,24 +36,3 @@ export class JsonKeyValueAdapter implements IKeyValueAdapter {
3436
return [setting.key, parsedValue];
3537
}
3638
}
37-
38-
// Determine whether a content type string is a valid JSON content type.
39-
// https://docs.microsoft.com/en-us/azure/azure-app-configuration/howto-leverage-json-content-type
40-
function isJsonContentType(contentTypeValue: string): boolean {
41-
if (!contentTypeValue) {
42-
return false;
43-
}
44-
45-
const contentTypeNormalized: string = contentTypeValue.trim().toLowerCase();
46-
const mimeType: string = contentTypeNormalized.split(";", 1)[0].trim();
47-
const typeParts: string[] = mimeType.split("/");
48-
if (typeParts.length !== 2) {
49-
return false;
50-
}
51-
52-
if (typeParts[0] !== "application") {
53-
return false;
54-
}
55-
56-
return typeParts[1].split("+").includes("json");
57-
}

‎src/common/contentType.ts‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
// Copyright (c) Microsoft Corporation.
2+
// Licensed under the MIT license.
3+
4+
export type ContentType = {
5+
mediaType: string;
6+
parameters: Record<string, string>;
7+
}
8+
9+
export function parseContentType(contentTypeValue: string | undefined): ContentType | undefined {
10+
if (!contentTypeValue) {
11+
return undefined;
12+
}
13+
const [mediaType, ...args] = contentTypeValue.split(";").map((s) => s.trim());
14+
const parameters: Record<string, string> = {};
15+
16+
for (const param of args) {
17+
const [key, value] = param.split("=").map((s) => s.trim());
18+
if (key && value) {
19+
parameters[key] = value;
20+
}
21+
}
22+
23+
return { mediaType, parameters };
24+
}
25+
26+
// Determine whether a content type string is a valid JSON content type.
27+
// https://docs.microsoft.com/en-us/azure/azure-app-configuration/howto-leverage-json-content-type
28+
export function isJsonContentType(contentType: ContentType | undefined): boolean {
29+
const mediaType = contentType?.mediaType?.trim().toLowerCase();
30+
if (!mediaType) {
31+
return false;
32+
}
33+
34+
const typeParts: string[] = mediaType.split("/");
35+
if (typeParts.length !== 2) {
36+
return false;
37+
}
38+
39+
if (typeParts[0] !== "application") {
40+
return false;
41+
}
42+
43+
return typeParts[1].split("+").includes("json");
44+
}

‎src/requestTracing/AIConfigurationTracingOptions.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,8 @@ export class AIConfigurationTracingOptions {
99
this.usesAIConfiguration = false;
1010
this.usesAIChatCompletionConfiguration = false;
1111
}
12+
13+
usesAnyTracingFeature() {
14+
return this.usesAIConfiguration || this.usesAIChatCompletionConfiguration;
15+
}
1216
}

‎src/requestTracing/utils.ts‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,10 @@ import {
2929
FAILOVER_REQUEST_TAG,
3030
FEATURES_KEY,
3131
LOAD_BALANCE_CONFIGURED_TAG,
32-
FM_VERSION_KEY
32+
FM_VERSION_KEY,
33+
DELIMITER,
34+
AI_CONFIGURATION_TAG,
35+
AI_CHAT_COMPLETION_CONFIGURATION_TAG
3336
} from "./constants";
3437

3538
export interface RequestTracingOptions {
@@ -127,9 +130,22 @@ export function createCorrelationContextHeader(requestTracingOptions: RequestTra
127130
keyValues.set(FM_VERSION_KEY, requestTracingOptions.fmVersion);
128131
}
129132

130-
// Compact tags: Features=LB+...
131-
if (appConfigOptions?.loadBalancingEnabled) {
132-
keyValues.set(FEATURES_KEY, LOAD_BALANCE_CONFIGURED_TAG);
133+
// Compact tags: Features=LB+AI+AICC...
134+
if (appConfigOptions?.loadBalancingEnabled || requestTracingOptions.aiConfigurationTracing?.usesAnyTracingFeature()) {
135+
const tags: string[] = [];
136+
if (appConfigOptions?.loadBalancingEnabled) {
137+
tags.push(LOAD_BALANCE_CONFIGURED_TAG);
138+
}
139+
if (requestTracingOptions.aiConfigurationTracing?.usesAIConfiguration) {
140+
tags.push(AI_CONFIGURATION_TAG);
141+
}
142+
if (requestTracingOptions.aiConfigurationTracing?.usesAIChatCompletionConfiguration) {
143+
tags.push(AI_CHAT_COMPLETION_CONFIGURATION_TAG);
144+
}
145+
146+
if (tags.length > 0) {
147+
keyValues.set(FEATURES_KEY, tags.join(DELIMITER));
148+
}
133149
}
134150

135151
const contextParts: string[] = [];

‎test/requestTracing.test.ts‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,19 @@ describe("request tracing", function () {
6464
expect(correlationContext.includes("UsesKeyVault")).eq(true);
6565
});
6666

67+
it("should have loadbalancing tag in correlation-context header", async () => {
68+
try {
69+
await load(createMockedConnectionString(fakeEndpoint), {
70+
clientOptions,
71+
loadBalancingEnabled: true,
72+
});
73+
} catch (e) { /* empty */ }
74+
expect(headerPolicy.headers).not.undefined;
75+
const correlationContext = headerPolicy.headers.get("Correlation-Context");
76+
expect(correlationContext).not.undefined;
77+
expect(correlationContext.includes("Features=LB")).eq(true);
78+
});
79+
6780
it("should have replica count in correlation-context header", async () => {
6881
const replicaCount = 2;
6982
sinon.stub(ConfigurationClientManager.prototype, "getReplicaCount").returns(replicaCount);
@@ -293,6 +306,37 @@ describe("request tracing", function () {
293306
restoreMocks();
294307
});
295308

309+
it("should have AI tag in correlation-context header if key values use AI configuration", async () => {
310+
let correlationContext: string = "";
311+
const listKvCallback = (listOptions) => {
312+
correlationContext = listOptions?.requestOptions?.customHeaders[CORRELATION_CONTEXT_HEADER_NAME] ?? "";
313+
};
314+
315+
mockAppConfigurationClientListConfigurationSettings([[
316+
createMockedKeyValue({ contentType: "application/json; profile=\"https://azconfig.io/mime-profiles/ai/chat-completion\"" })
317+
]], listKvCallback);
318+
319+
const settings = await load(createMockedConnectionString(fakeEndpoint), {
320+
refreshOptions: {
321+
enabled: true,
322+
refreshIntervalInMs: 1000
323+
}
324+
});
325+
326+
expect(correlationContext).not.undefined;
327+
expect(correlationContext?.includes("RequestType=Startup")).eq(true);
328+
329+
await sleepInMs(1000 + 1);
330+
try {
331+
await settings.refresh();
332+
} catch (e) { /* empty */ }
333+
expect(headerPolicy.headers).not.undefined;
334+
expect(correlationContext).not.undefined;
335+
expect(correlationContext?.includes("Features=AI+AICC")).eq(true);
336+
337+
restoreMocks();
338+
});
339+
296340
describe("request tracing in Web Worker environment", () => {
297341
let originalNavigator;
298342
let originalWorkerNavigator;

0 commit comments

Comments
 (0)