Skip to content

Commit a17ad43

Browse files
yasithrashanyasithrashan
committed
Implement AI telemetry event tracking
1 parent fb28682 commit a17ad43

File tree

8 files changed

+132
-15
lines changed

8 files changed

+132
-15
lines changed

‎workspaces/ballerina/ballerina-extension/src/extension.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ export async function activate(context: ExtensionContext) {
110110
extension.context = context;
111111
// Init RPC Layer methods
112112
RPCLayer.init();
113-
113+
114114
// Wait for the ballerina extension to be ready
115115
await StateMachine.initialize();
116116

‎workspaces/ballerina/ballerina-extension/src/features/ai/service/design/design.ts‎

Lines changed: 74 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
// specific language governing permissions and limitations
1515
// under the License.
1616

17-
import { Command, GenerateAgentCodeRequest, ProjectSource, AIChatMachineEventType} from "@wso2/ballerina-core";
17+
import { Command, GenerateAgentCodeRequest, ProjectSource, AIChatMachineEventType } from "@wso2/ballerina-core";
1818
import { ModelMessage, stepCountIs, streamText } from "ai";
1919
import { getAnthropicClient, getProviderCacheControl, ANTHROPIC_SONNET_4 } from "../connection";
2020
import { getErrorMessage, populateHistoryForAgent } from "../utils";
@@ -38,6 +38,10 @@ import { LangfuseExporter } from 'langfuse-vercel';
3838
import { NodeSDK } from '@opentelemetry/sdk-node';
3939
import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
4040
import { getProjectSource } from "../../utils/project-utils";
41+
import { sendTelemetryEvent, sendTelemetryException } from "../../../telemetry";
42+
import { TM_EVENT_BALLERINA_AI_GENERATION_SUBMITTED, TM_EVENT_BALLERINA_AI_GENERATION_COMPLETED, TM_EVENT_BALLERINA_AI_GENERATION_FAILED, TM_EVENT_BALLERINA_AI_GENERATION_ABORTED, TM_EVENT_BALLERINA_AI_GENERATION_DIAGNOSTICS, CMP_BALLERINA_AI_GENERATION } from "../../../telemetry";
43+
import { extension } from "../../../../BalExtensionContext";
44+
4145

4246
const LANGFUSE_SECRET = process.env.LANGFUSE_SECRET;
4347
const LANGFUSE_PUBLIC = process.env.LANGFUSE_PUBLIC;
@@ -74,6 +78,19 @@ export async function generateDesignCore(
7478

7579
const cacheOptions = await getProviderCacheControl();
7680

81+
// Get state machine context for telemetry
82+
const stateContext = AIChatStateMachine.context();
83+
84+
// Send telemetry when the user submits a query
85+
sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_SUBMITTED, CMP_BALLERINA_AI_GENERATION, {
86+
projectId: stateContext.projectId || 'unknown',
87+
messageId: messageId,
88+
command: Command.Design,
89+
operationType: params.operationType,
90+
isPlanMode: isPlanModeEnabled.toString(),
91+
approvalMode: stateContext.autoApproveEnabled ? 'auto' : 'manual',
92+
});
93+
7794
const modifiedFiles: string[] = [];
7895

7996
const userMessageContent = getUserPrompt(params.usecase, tempProjectPath, projects, isPlanModeEnabled, params.codeContext);
@@ -108,6 +125,9 @@ export async function generateDesignCore(
108125
[DIAGNOSTICS_TOOL_NAME]: createDiagnosticsTool(tempProjectPath),
109126
};
110127

128+
// Timing metrics for telemetry - capture at generation start
129+
const generationStartTime = Date.now();
130+
111131
const { fullStream, response } = streamText({
112132
model: await getAnthropicClient(ANTHROPIC_SONNET_4),
113133
maxOutputTokens: 8192,
@@ -140,6 +160,7 @@ export async function generateDesignCore(
140160
break;
141161
}
142162
case "tool-call": {
163+
143164
const toolName = part.toolName;
144165
accumulateToolCall(currentAssistantContent, part);
145166

@@ -227,6 +248,22 @@ export async function generateDesignCore(
227248
if (shouldCleanup) {
228249
cleanupTempProject(tempProjectPath);
229250
}
251+
252+
// Send telemetry for generation error
253+
const errorObj = error instanceof Error ? error : new Error(String(error));
254+
const errorTime = Date.now();
255+
sendTelemetryException(extension.ballerinaExtInstance, errorObj, CMP_BALLERINA_AI_GENERATION, {
256+
event: TM_EVENT_BALLERINA_AI_GENERATION_FAILED,
257+
projectId: stateContext.projectId || 'unknown',
258+
messageId: messageId,
259+
errorMessage: getErrorMessage(error),
260+
errorType: errorObj.name || 'Unknown',
261+
generationStartTime: generationStartTime.toString(),
262+
errorTime: errorTime.toString(),
263+
durationMs: (errorTime - generationStartTime).toString(),
264+
});
265+
266+
230267
eventHandler({ type: "error", content: getErrorMessage(error) });
231268
return tempProjectPath;
232269
}
@@ -236,6 +273,18 @@ export async function generateDesignCore(
236273
}
237274
case "abort": {
238275
console.log("[Design] Aborted by user.");
276+
const abortTime = Date.now();
277+
278+
// Send telemetry for generation abort
279+
sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_ABORTED, CMP_BALLERINA_AI_GENERATION, {
280+
projectId: stateContext.projectId || 'unknown',
281+
messageId: messageId,
282+
generationStartTime: generationStartTime.toString(),
283+
abortTime: abortTime.toString(),
284+
durationMs: (abortTime - generationStartTime).toString(),
285+
modifiedFilesCount: modifiedFiles.length.toString(),
286+
});
287+
239288
let messagesToSave: any[] = [];
240289
try {
241290
const partialResponse = await response;
@@ -277,9 +326,19 @@ Generation stopped by user. The last in-progress task was not saved. Files have
277326
case "finish": {
278327
const finalResponse = await response;
279328
const assistantMessages = finalResponse.messages || [];
329+
const generationEndTime = Date.now();
280330

281331
const finalDiagnostics = await checkCompilationErrors(tempProjectPath);
282332
if (finalDiagnostics.diagnostics && finalDiagnostics.diagnostics.length > 0) {
333+
// Send telemetry for final diagnostics check
334+
sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_DIAGNOSTICS, CMP_BALLERINA_AI_GENERATION, {
335+
projectId: stateContext.projectId || 'unknown',
336+
messageId: messageId,
337+
hasErrors: 'true',
338+
errorCount: finalDiagnostics.diagnostics.length.toString(),
339+
stage: 'completion',
340+
});
341+
283342
eventHandler({
284343
type: "diagnostics",
285344
diagnostics: finalDiagnostics.diagnostics
@@ -297,14 +356,27 @@ Generation stopped by user. The last in-progress task was not saved. Files have
297356

298357
updateAndSaveChat(messageId, userMessageContent, assistantMessages, eventHandler);
299358
eventHandler({ type: "stop", command: Command.Design });
359+
360+
// Send telemetry for generation completion
361+
sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_COMPLETED, CMP_BALLERINA_AI_GENERATION, {
362+
projectId: stateContext.projectId || 'unknown',
363+
messageId: messageId,
364+
modifiedFilesCount: modifiedFiles.length.toString(),
365+
generationStartTime: generationStartTime.toString(),
366+
generationEndTime: generationEndTime.toString(),
367+
durationMs: (generationEndTime - generationStartTime).toString(),
368+
isPlanMode: isPlanModeEnabled.toString(),
369+
approvalMode: stateContext.autoApproveEnabled ? 'auto' : 'manual',
370+
});
371+
300372
AIChatStateMachine.sendEvent({
301373
type: AIChatMachineEventType.FINISH_EXECUTION,
302374
});
303375
await langfuseExporter.forceFlush();
304376
return tempProjectPath;
305377
}
306378
}
307-
}
379+
}
308380

309381
return tempProjectPath;
310382
}

‎workspaces/ballerina/ballerina-extension/src/features/telemetry/activator.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import { debug } from "../../utils";
2121
import { window } from "vscode";
2222
import {
2323
CMP_EDITOR_SUPPORT, getMessageObject, getTelemetryProperties, sendTelemetryEvent, TM_ERROR_LANG_SERVER,
24-
TM_EVENT_EDIT_DIAGRAM, TM_EVENT_EDIT_SOURCE, TM_EVENT_KILL_TERMINAL, TM_FEATURE_USAGE_LANG_SERVER
24+
TM_EVENT_EDIT_DIAGRAM, TM_EVENT_EDIT_SOURCE, TM_EVENT_KILL_TERMINAL, TM_FEATURE_USAGE_LANG_SERVER,
2525
} from ".";
2626

2727
const schedule = require('node-schedule');
@@ -35,20 +35,20 @@ export function activate(ballerinaExtInstance: BallerinaExtension) {
3535
const langClient = <ExtendedLangClient>ballerinaExtInstance.langClient;
3636

3737
// Start listening telemtry events from language server
38-
langClient.onNotification('telemetry/event', (event: LSTelemetryEvent) => {
38+
langClient.onNotification('telemetry/event', async (event: LSTelemetryEvent) => {
3939
let props: { [key: string]: string; };
4040
switch (event.type) {
4141
case TM_EVENT_TYPE_ERROR:
4242
const errorEvent: LSErrorTelemetryEvent = <LSErrorTelemetryEvent>event;
43-
props = getTelemetryProperties(ballerinaExtInstance, event.component, getMessageObject(TM_EVENT_TYPE_ERROR));
43+
props = await getTelemetryProperties(ballerinaExtInstance, event.component, getMessageObject(TM_EVENT_TYPE_ERROR));
4444
props["ballerina.langserver.error.description"] = errorEvent.message;
4545
props["ballerina.langserver.error.stacktrace"] = errorEvent.errorStackTrace;
4646
props["ballerina.langserver.error.message"] = errorEvent.errorMessage;
4747
reporter.sendTelemetryEvent(TM_ERROR_LANG_SERVER, props);
4848
break;
4949
case TM_EVENT_TYPE_FEATURE_USAGE:
5050
const usageEvent: LSFeatureUsageTelemetryEvent = <LSFeatureUsageTelemetryEvent>event;
51-
props = getTelemetryProperties(ballerinaExtInstance, event.component,
51+
props = await getTelemetryProperties(ballerinaExtInstance, event.component,
5252
getMessageObject(TM_EVENT_TYPE_FEATURE_USAGE));
5353
props["ballerina.langserver.feature.name"] = usageEvent.featureName;
5454
props["ballerina.langserver.feature.class"] = usageEvent.featureClass;

‎workspaces/ballerina/ballerina-extension/src/features/telemetry/components.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,3 +45,5 @@ export const CMP_CHOREO_AUTHENTICATION = "component.choreo.authentication";
4545
export const CMP_PERF_ANALYZER = "component.perf.analyzer";
4646
export const CMP_NOTEBOOK = "component.notebook";
4747
export const CMP_OPEN_VSCODE_URL = "component.open.vscode.url";
48+
49+
export const CMP_BALLERINA_AI_GENERATION = "ballerina.ai.generation";

‎workspaces/ballerina/ballerina-extension/src/features/telemetry/events.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,3 +122,11 @@ export const TM_EVENT_OPEN_REPO_CHANGE_PATH = "vscode.open.repo.change.path";
122122
export const TM_EVENT_OPEN_REPO_CANCELED = "vscode.open.repo.canceled";
123123
export const TM_EVENT_OPEN_REPO_NEW_FOLDER = "vscode.open.exist.repo.new.folder";
124124
export const TM_EVENT_OPEN_REPO_SAME_FOLDER = "vscode.open.exist.repo.same.folder";
125+
126+
// events for AI features
127+
export const TM_EVENT_BALLERINA_AI_GENERATION_SUBMITTED = "ballerina.ai.generation.submitted";
128+
export const TM_EVENT_BALLERINA_AI_GENERATION_COMPLETED = "ballerina.ai.generation.completed";
129+
export const TM_EVENT_BALLERINA_AI_GENERATION_FAILED = "ballerina.ai.generation.failed";
130+
export const TM_EVENT_BALLERINA_AI_GENERATION_ABORTED = "ballerina.ai.generation.aborted";
131+
export const TM_EVENT_BALLERINA_AI_GENERATION_DIAGNOSTICS = "ballerina.ai.generation.diagnostics";
132+
export const TM_EVENT_BALLERINA_AI_GENERATION_REVERTED = "ballerina.ai.generation.reverted";

‎workspaces/ballerina/ballerina-extension/src/features/telemetry/index.ts‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import TelemetryReporter from "vscode-extension-telemetry";
2020
import { BallerinaExtension } from "../../core";
21+
import { getLoginMethod, getBiIntelId } from "../../utils/ai/auth";
2122

2223
//Ballerina-VSCode-Extention repo key as default
2324
const DEFAULT_KEY = "3a82b093-5b7b-440c-9aa2-3b8e8e5704e7";
@@ -37,26 +38,30 @@ export function createTelemetryReporter(ext: BallerinaExtension): TelemetryRepor
3738
return reporter;
3839
}
3940

40-
export function sendTelemetryEvent(extension: BallerinaExtension, eventName: string, componentName: string,
41+
export async function sendTelemetryEvent(extension: BallerinaExtension, eventName: string, componentName: string,
4142
customDimensions: { [key: string]: string; } = {}, measurements: { [key: string]: number; } = {}) {
4243
// temporarily disabled in codeserver due to GDPR issue
4344
if (extension.isTelemetryEnabled() && !extension.getCodeServerContext().codeServerEnv) {
44-
extension.telemetryReporter.sendTelemetryEvent(eventName, getTelemetryProperties(extension, componentName,
45+
extension.telemetryReporter.sendTelemetryEvent(eventName, await getTelemetryProperties(extension, componentName,
4546
customDimensions), measurements);
4647
}
4748
}
4849

49-
export function sendTelemetryException(extension: BallerinaExtension, error: Error, componentName: string,
50+
export async function sendTelemetryException(extension: BallerinaExtension, error: Error, componentName: string,
5051
params: { [key: string]: string } = {}) {
5152
// temporarily disabled in codeserver due to GDPR issue
5253
if (extension.isTelemetryEnabled() && !extension.getCodeServerContext().codeServerEnv) {
53-
extension.telemetryReporter.sendTelemetryException(error, getTelemetryProperties(extension, componentName,
54+
extension.telemetryReporter.sendTelemetryException(error, await getTelemetryProperties(extension, componentName,
5455
params));
5556
}
5657
}
5758

58-
export function getTelemetryProperties(extension: BallerinaExtension, component: string, params: { [key: string]: string; } = {})
59-
: { [key: string]: string; } {
59+
export async function getTelemetryProperties(extension: BallerinaExtension, component: string, params: { [key: string]: string; } = {})
60+
: Promise<{ [key: string]: string; }> {
61+
62+
const userType = await getLoginMethod();
63+
const biIntelId = await getBiIntelId();
64+
6065
return {
6166
...params,
6267
'ballerina.version': extension ? extension.ballerinaVersion : '',
@@ -69,9 +74,10 @@ export function getTelemetryProperties(extension: BallerinaExtension, component:
6974
'component': CHOREO_COMPONENT_ID,
7075
'project': CHOREO_PROJECT_ID,
7176
'org': CHOREO_ORG_ID,
77+
'userType': userType,
78+
'biIntelId': biIntelId,
7279
};
7380
}
74-
7581
export function getMessageObject(message?: string): { [key: string]: string; } {
7682
if (message) {
7783
return { 'ballerina.message': message };

‎workspaces/ballerina/ballerina-extension/src/utils/ai/auth.ts‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ vscode.authentication.onDidChangeSessions(async e => {
100100
await extension.context.secrets.delete('GITHUB_COPILOT_TOKEN');
101101
await extension.context.secrets.delete('GITHUB_TOKEN');
102102
} else {
103-
//it could be a login(which we havent captured) or a logout
103+
//it could be a login(which we havent captured) or a logout
104104
// vscode.window.showInformationMessage(
105105
// 'WSO2 Integrator: BI supports completions with GitHub Copilot.',
106106
// 'Login with GitHub Copilot'
@@ -229,6 +229,25 @@ export const getAwsBedrockCredentials = async (): Promise<{
229229
return credentials.secrets;
230230
};
231231

232+
// ==================================
233+
// Unique user identifier for BIIntel
234+
// ==================================
235+
export const getBiIntelId = async (): Promise<string | undefined> => {
236+
try {
237+
const credentials = await getAuthCredentials();
238+
if (!credentials || credentials.loginMethod !== LoginMethod.BI_INTEL) {
239+
return undefined;
240+
}
241+
242+
const { accessToken } = credentials.secrets;
243+
const decoded = jwtDecode<JwtPayload>(accessToken);
244+
return decoded.sub;
245+
} catch (error) {
246+
console.error('Error decoding JWT token:', error);
247+
return undefined;
248+
}
249+
};
250+
232251
export const getRefreshedAccessToken = async (): Promise<string> => {
233252
return new Promise(async (resolve, reject) => {
234253
const CommonReqHeaders = {

‎workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,10 @@ import { addUserMessage, updateChatMessage, convertChatHistoryToModelMessages, c
3333
import { saveChatState, loadChatState, clearChatState, clearChatStateAction, getAllProjectIds, clearAllChatStates, getChatStateMetadata } from './chatStatePersistence';
3434
import { normalizeCodeContext } from './codeContextUtils';
3535

36+
import { sendTelemetryEvent } from '../../features/telemetry';
37+
import { TM_EVENT_BALLERINA_AI_GENERATION_REVERTED, CMP_BALLERINA_AI_GENERATION } from '../../features/telemetry';
38+
39+
3640
const cleanupOldCheckpoints = (checkpoints: Checkpoint[]): Checkpoint[] => {
3741
const config = getCheckpointConfig();
3842
if (checkpoints.length <= config.maxCount) {
@@ -88,6 +92,12 @@ const restoreCheckpointAction = (context: AIChatMachineContext, event: any) => {
8892
return;
8993
}
9094

95+
// Send telemetry when the user clicks the revert button
96+
sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_REVERTED, CMP_BALLERINA_AI_GENERATION, {
97+
messageId: checkpoint.messageId,
98+
checkpointId: checkpointId,
99+
});
100+
91101
const messageIndex = context.chatHistory.findIndex(m => m.id === checkpoint.messageId);
92102
const restoredHistory = messageIndex >= 0 ? context.chatHistory.slice(0, messageIndex) : context.chatHistory;
93103

0 commit comments

Comments
 (0)