Enhance AI generation telemetry tracking

yasithrashan · yasithrashan · commit df75f81cc477 · 2025-12-03T22:09:53.000+05:30
diff --git a/workspaces/ballerina/ballerina-extension/src/extension.ts b/workspaces/ballerina/ballerina-extension/src/extension.ts
@@ -113,8 +113,8 @@ export async function activate(context: ExtensionContext) {
 
     // Wait for the ballerina extension to be ready
     await StateMachine.initialize();
-
     // Then return the ballerina extension context
+    
     return { ballerinaExtInstance: extension.ballerinaExtInstance, projectPath: StateMachine.context().projectPath };
 }
 
diff --git a/workspaces/ballerina/ballerina-extension/src/features/ai/service/design/design.ts b/workspaces/ballerina/ballerina-extension/src/features/ai/service/design/design.ts
@@ -35,7 +35,7 @@ import { createConnectorGeneratorTool, CONNECTOR_GENERATOR_TOOL } from "../libs/
 import { LangfuseExporter } from 'langfuse-vercel';
 import { NodeSDK } from '@opentelemetry/sdk-node';
 import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
-import { sendTelemetryEvent, TM_EVENT_BALLERINA_AI_QUERY_SUBMIT, TM_EVENT_BALLERINA_AI_GENERATION_FINISHED, CMP_BALLERINA_AI } from "../../../telemetry";
+import { sendTelemetryEvent, sendTelemetryException, TM_EVENT_BALLERINA_AI_QUERY_SUBMIT, TM_EVENT_BALLERINA_AI_GENERATION_COMPLETED, TM_EVENT_BALLERINA_AI_GENERATION_ERROR, TM_EVENT_BALLERINA_AI_GENERATION_ABORTED, TM_EVENT_BALLERINA_AI_DIAGNOSTICS, CMP_BALLERINA_AI } from "../../../telemetry";
 import { extension } from "../../../../BalExtensionContext";
 
 const LANGFUSE_SECRET = process.env.LANGFUSE_SECRET;
@@ -64,11 +64,17 @@ export async function generateDesignCore(params: GenerateAgentCodeRequest, event
 
     const cacheOptions = await getProviderCacheControl();
 
+    // Get state machine context for telemetry
+    const stateContext = AIChatStateMachine.context();
+
     // Send telemetry when the user submits a query
     sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_QUERY_SUBMIT, CMP_BALLERINA_AI, {
+        projectId: stateContext.projectId || 'unknown',
         messageId: messageId,
         command: Command.Design,
         operationType: params.operationType,
+        isPlanMode: isPlanModeEnabled.toString(),
+        approvalMode: stateContext.autoApproveEnabled ? 'auto' : 'manual',
     });
 
     const modifiedFiles: string[] = [];
@@ -127,15 +133,36 @@ export async function generateDesignCore(params: GenerateAgentCodeRequest, event
     let accumulatedMessages: any[] = [];
     let currentAssistantContent: any[] = [];
 
+    // Timing metrics for telemetry
+    let generationStartTime: number | undefined;
+    let firstTokenTime: number | undefined;
+    let lastTokenTime: number | undefined;
+
     for await (const part of fullStream) {
         switch (part.type) {
             case "text-delta": {
+                // Capture first token time
+                if (!firstTokenTime) {
+                    firstTokenTime = Date.now();
+                    generationStartTime = firstTokenTime;
+                }
+                // Update last token time
+                lastTokenTime = Date.now();
+
                 const textPart = part.text;
                 eventHandler({ type: "content_block", content: textPart });
                 accumulateTextContent(currentAssistantContent, textPart);
                 break;
             }
             case "tool-call": {
+                // Capture first token time for tool calls as well
+                if (!firstTokenTime) {
+                    firstTokenTime = Date.now();
+                    generationStartTime = firstTokenTime;
+                }
+                // Update last token time
+                lastTokenTime = Date.now();
+
                 const toolName = part.toolName;
                 accumulateToolCall(currentAssistantContent, part);
 
@@ -207,6 +234,20 @@ export async function generateDesignCore(params: GenerateAgentCodeRequest, event
                         toolOutput: { success: true, action }
                     });
                 } else if (toolName === DIAGNOSTICS_TOOL_NAME) {
+                    // Send telemetry for diagnostics
+                    const diagnosticsResult = result as any;
+                    const hasErrors = diagnosticsResult?.diagnostics?.some((d: any) => d.severity === 'Error') || false;
+                    const errorCount = diagnosticsResult?.diagnostics?.filter((d: any) => d.severity === 'Error').length || 0;
+                    const warningCount = diagnosticsResult?.diagnostics?.filter((d: any) => d.severity === 'Warning').length || 0;
+
+                    sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_DIAGNOSTICS, CMP_BALLERINA_AI, {
+                        projectId: stateContext.projectId || 'unknown',
+                        messageId: messageId,
+                        hasErrors: hasErrors.toString(),
+                        errorCount: errorCount.toString(),
+                        warningCount: warningCount.toString(),
+                    });
+
                     eventHandler({
                         type: "tool_result",
                         toolName,
@@ -220,8 +261,26 @@ export async function generateDesignCore(params: GenerateAgentCodeRequest, event
             case "error": {
                 const error = part.error;
                 console.error("[Design] Error:", error);
+
+                // Send telemetry for generation error
+                const errorObj = error instanceof Error ? error : new Error(String(error));
+                sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_ERROR, CMP_BALLERINA_AI, {
+                    projectId: stateContext.projectId || 'unknown',
+                    messageId: messageId,
+                    errorMessage: getErrorMessage(error),
+                    errorType: errorObj.name || 'Unknown',
+                    generationStartTime: generationStartTime?.toString() || 'unknown',
+                    errorTime: Date.now().toString(),
+                });
+
+                sendTelemetryException(extension.ballerinaExtInstance, errorObj, CMP_BALLERINA_AI, {
+                    projectId: stateContext.projectId || 'unknown',
+                    messageId: messageId,
+                });
+
                 // Cleanup temp project on error
                 cleanupTempProject(tempProjectPath);
+
                 eventHandler({ type: "error", content: getErrorMessage(error) });
                 break;
             }
@@ -231,6 +290,19 @@ export async function generateDesignCore(params: GenerateAgentCodeRequest, event
             }
             case "abort": {
                 console.log("[Design] Aborted by user.");
+                const abortTime = Date.now();
+
+                // Send telemetry for generation abort
+                sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_ABORTED, CMP_BALLERINA_AI, {
+                    projectId: stateContext.projectId || 'unknown',
+                    messageId: messageId,
+                    generationStartTime: generationStartTime?.toString() || 'unknown',
+                    abortTime: abortTime.toString(),
+                    firstTokenTime: firstTokenTime?.toString() || 'unknown',
+                    lastTokenTime: lastTokenTime?.toString() || 'unknown',
+                    modifiedFilesCount: modifiedFiles.length.toString(),
+                });
+
                 let messagesToSave: any[] = [];
                 try {
                     const partialResponse = await response;
@@ -272,6 +344,7 @@ Generation stopped by user. The last in-progress task was not saved. Files have
                 const finishReason = part.finishReason;
                 const finalResponse = await response;
                 const assistantMessages = finalResponse.messages || [];
+                const generationEndTime = Date.now();
 
                 console.log(`[Design] Finished with reason: ${finishReason}`);
 
@@ -288,11 +361,17 @@ Generation stopped by user. The last in-progress task was not saved. Files have
                 updateAndSaveChat(messageId, userMessageContent, assistantMessages, eventHandler);
                 eventHandler({ type: "stop", command: Command.Design });
 
-                // Send telemetry after generation
-                sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_FINISHED, CMP_BALLERINA_AI, {
+                sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_COMPLETED, CMP_BALLERINA_AI, {
+                    projectId: stateContext.projectId || 'unknown',
                     messageId: messageId,
                     finishReason: finishReason,
                     modifiedFilesCount: modifiedFiles.length.toString(),
+                    generationStartTime: generationStartTime?.toString() || 'unknown',
+                    generationEndTime: generationEndTime.toString(),
+                    firstTokenTime: firstTokenTime?.toString() || 'unknown',
+                    lastTokenTime: lastTokenTime?.toString() || 'unknown',
+                    isPlanMode: isPlanModeEnabled.toString(),
+                    approvalMode: stateContext.autoApproveEnabled ? 'auto' : 'manual',
                 });
 
                 AIChatStateMachine.sendEvent({
diff --git a/workspaces/ballerina/ballerina-extension/src/features/telemetry/events.ts b/workspaces/ballerina/ballerina-extension/src/features/telemetry/events.ts
@@ -125,5 +125,8 @@ export const TM_EVENT_OPEN_REPO_SAME_FOLDER = "vscode.open.exist.repo.same.folde
 
 // events for AI features
 export const TM_EVENT_BALLERINA_AI_QUERY_SUBMIT = "ballerina.ai.query.submit";
-export const TM_EVENT_BALLERINA_AI_GENERATION_FINISHED = "ballerina.ai.generation.finished";
+export const TM_EVENT_BALLERINA_AI_GENERATION_COMPLETED = "ballerina.ai.generation.completed";
+export const TM_EVENT_BALLERINA_AI_GENERATION_ERROR = "ballerina.ai.generation.error";
+export const TM_EVENT_BALLERINA_AI_GENERATION_ABORTED = "ballerina.ai.generation.aborted";
+export const TM_EVENT_BALLERINA_AI_DIAGNOSTICS = "ballerina.ai.diagnostics";
 export const TM_EVENT_BALLERINA_AI_REVERT = "ballerina.ai.code.revert";
diff --git a/workspaces/ballerina/ballerina-extension/src/features/telemetry/index.ts b/workspaces/ballerina/ballerina-extension/src/features/telemetry/index.ts
@@ -21,7 +21,8 @@ import { BallerinaExtension } from "../../core";
 import { getLoginMethod, getBiIntelId } from "../../utils/ai/auth";
 
 //Ballerina-VSCode-Extention repo key as default
-const DEFAULT_KEY = "3a82b093-5b7b-440c-9aa2-3b8e8e5704e7";
+// const DEFAULT_KEY = "3a82b093-5b7b-440c-9aa2-3b8e8e5704e7";
+const DEFAULT_KEY = "ff7807cf-db47-4e8b-bd70-03fdfc2576b7";
 const INSTRUMENTATION_KEY = process.env.CODE_SERVER_ENV && process.env.VSCODE_CHOREO_INSTRUMENTATION_KEY ? process.env.VSCODE_CHOREO_INSTRUMENTATION_KEY : DEFAULT_KEY;
 const isWSO2User = process.env.VSCODE_CHOREO_USER_EMAIL ? process.env.VSCODE_CHOREO_USER_EMAIL.endsWith('@wso2.com') : false;
 const isAnonymous = process.env.VSCODE_CHOREO_USER_EMAIL ? process.env.VSCODE_CHOREO_USER_EMAIL.endsWith('@choreo.dev') : false;
@@ -59,7 +60,7 @@ export async function sendTelemetryException(extension: BallerinaExtension, erro
 export async function getTelemetryProperties(extension: BallerinaExtension, component: string, params: { [key: string]: string; } = {})
     : Promise<{ [key: string]: string; }> {
 
-    const loginType = await getLoginMethod();
+    const userType = await getLoginMethod();
     const biIntelId = await getBiIntelId();
 
     return {
@@ -74,7 +75,7 @@ export async function getTelemetryProperties(extension: BallerinaExtension, comp
         'component': CHOREO_COMPONENT_ID,
         'project': CHOREO_PROJECT_ID,
         'org': CHOREO_ORG_ID,
-        'loginType': loginType,
+        'userType': userType,
         'biIntelId': biIntelId,
     }
 }
diff --git a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/aiChatMachine.ts
@@ -25,8 +25,8 @@ import { generateDesign } from '../../features/ai/service/design/design';
 import { captureWorkspaceSnapshot, restoreWorkspaceSnapshot } from './checkpoint/checkpointUtils';
 import { getCheckpointConfig } from './checkpoint/checkpointConfig';
 import { notifyCheckpointCaptured } from '../../RPCLayer';
+import { StateMachine } from '../../stateMachine';
 import { sendTelemetryEvent, TM_EVENT_BALLERINA_AI_REVERT, CMP_BALLERINA_AI } from '../../features/telemetry';
-
 // Extracted utilities
 import { generateProjectId, generateSessionId } from './idGenerators';
 import { addUserMessage, updateChatMessage, convertChatHistoryToModelMessages, convertChatHistoryToUIMessages } from './chatHistoryUtils';

Original file line number	Diff line number	Diff line change
`@@ -113,8 +113,8 @@ export async function activate(context: ExtensionContext) {`
`113`	`113`
`114`	`114`	`// Wait for the ballerina extension to be ready`
`115`	`115`	`await StateMachine.initialize();`
`116`		`-`
`117`	`116`	`// Then return the ballerina extension context`
	`117`	`+`
`118`	`118`	`return { ballerinaExtInstance: extension.ballerinaExtInstance, projectPath: StateMachine.context().projectPath };`
`119`	`119`	`}`
`120`	`120`