Implement telemetry tracking for AI features

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

@@ -110,7 +110,7 @@ export async function activate(context: ExtensionContext) {
     extension.context = context;
     // Init RPC Layer methods
     RPCLayer.init();
-    
+
     // Wait for the ballerina extension to be ready
     await StateMachine.initialize();
 

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

@@ -14,7 +14,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-import { Command, GenerateAgentCodeRequest, ProjectSource, AIChatMachineEventType} from "@wso2/ballerina-core";
+import { Command, GenerateAgentCodeRequest, ProjectSource, AIChatMachineEventType } from "@wso2/ballerina-core";
 import { ModelMessage, stepCountIs, streamText } from "ai";
 import { getAnthropicClient, getProviderCacheControl, ANTHROPIC_SONNET_4 } from "../connection";
 import { getErrorMessage, populateHistoryForAgent } from "../utils";
@@ -38,6 +38,11 @@ import { LangfuseExporter } from 'langfuse-vercel';
 import { NodeSDK } from '@opentelemetry/sdk-node';
 import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
 import { getProjectSource } from "../../utils/project-utils";
+import { sendTelemetryEvent, sendTelemetryException } from "../../../telemetry";
+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, CMP_BALLERINA_AI_GENERATION } from "../../../telemetry";
+import { extension } from "../../../../BalExtensionContext";
+import { getProjectMetrics } from "../../../telemetry/common/project-metrics";
+
 
 const LANGFUSE_SECRET = process.env.LANGFUSE_SECRET;
 const LANGFUSE_PUBLIC = process.env.LANGFUSE_PUBLIC;
@@ -74,6 +79,21 @@ export async function generateDesignCore(
 
     const cacheOptions = await getProviderCacheControl();
 
+    const stateContext = AIChatStateMachine.context();
+    const projectMetrics = await getProjectMetrics();
+
+    // Send telemetry when the user submits a query
+    sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_SUBMITTED, CMP_BALLERINA_AI_GENERATION, {
+        projectId: stateContext.projectId || 'unknown',
+        messageId: messageId,
+        command: Command.Design,
+        operationType: params.operationType,
+        isPlanMode: isPlanModeEnabled.toString(),
+        approvalMode: stateContext.autoApproveEnabled ? 'auto' : 'manual',
+        inputFileCount: projectMetrics.fileCount.toString(),
+        inputLineCount: projectMetrics.lineCount.toString(),
+    });
+
     const modifiedFiles: string[] = [];
 
     const userMessageContent = getUserPrompt(params.usecase, tempProjectPath, projects, isPlanModeEnabled, params.codeContext);
@@ -108,7 +128,14 @@ export async function generateDesignCore(
         [DIAGNOSTICS_TOOL_NAME]: createDiagnosticsTool(tempProjectPath),
     };
 
-    const { fullStream, response } = streamText({
+    // Timing metrics for telemetry - capture at generation start
+    const generationStartTime = Date.now();
+
+    // Diagnostic tracking during generation - tracks all compilation errors that occur
+    let diagnosticCheckCount = 0;
+    let totalCompilationErrorsDuringGeneration = 0;
+
+    const { fullStream, response, usage } = streamText({
         model: await getAnthropicClient(ANTHROPIC_SONNET_4),
         maxOutputTokens: 8192,
         temperature: 0,
@@ -140,6 +167,7 @@ export async function generateDesignCore(
                 break;
             }
             case "tool-call": {
+
                 const toolName = part.toolName;
                 accumulateToolCall(currentAssistantContent, part);
 
@@ -211,6 +239,14 @@ export async function generateDesignCore(
                         toolOutput: { success: true, action }
                     });
                 } else if (toolName === DIAGNOSTICS_TOOL_NAME) {
+                    // Track diagnostic errors during generation
+                    const diagnosticsResult = result as any;
+                    if (diagnosticsResult && diagnosticsResult.diagnostics) {
+                        const errorCount = diagnosticsResult.diagnostics.length;
+                        diagnosticCheckCount++;
+                        totalCompilationErrorsDuringGeneration += errorCount;
+                    }
+
                     eventHandler({
                         type: "tool_result",
                         toolName,
@@ -227,6 +263,22 @@ export async function generateDesignCore(
                 if (shouldCleanup) {
                     cleanupTempProject(tempProjectPath);
                 }
+
+                // Send telemetry for generation error
+                const errorObj = error instanceof Error ? error : new Error(String(error));
+                const errorTime = Date.now();
+                sendTelemetryException(extension.ballerinaExtInstance, errorObj, CMP_BALLERINA_AI_GENERATION, {
+                    event: TM_EVENT_BALLERINA_AI_GENERATION_FAILED,
+                    projectId: stateContext.projectId || 'unknown',
+                    messageId: messageId,
+                    errorMessage: getErrorMessage(error),
+                    errorType: errorObj.name || 'Unknown',
+                    generationStartTime: generationStartTime.toString(),
+                    errorTime: errorTime.toString(),
+                    durationMs: (errorTime - generationStartTime).toString(),
+                });
+
+
                 eventHandler({ type: "error", content: getErrorMessage(error) });
                 return tempProjectPath;
             }
@@ -236,6 +288,18 @@ export async function generateDesignCore(
             }
             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_GENERATION, {
+                    projectId: stateContext.projectId || 'unknown',
+                    messageId: messageId,
+                    generationStartTime: generationStartTime.toString(),
+                    abortTime: abortTime.toString(),
+                    durationMs: (abortTime - generationStartTime).toString(),
+                    modifiedFilesCount: modifiedFiles.length.toString(),
+                });
+
                 let messagesToSave: any[] = [];
                 try {
                     const partialResponse = await response;
@@ -277,6 +341,12 @@ Generation stopped by user. The last in-progress task was not saved. Files have
             case "finish": {
                 const finalResponse = await response;
                 const assistantMessages = finalResponse.messages || [];
+                const generationEndTime = Date.now();
+
+                const usageInfo = await usage;
+                const inputTokens = usageInfo.inputTokens || 0;
+                const outputTokens = usageInfo.outputTokens || 0;
+                const totalTokens = usageInfo.totalTokens || 0;
 
                 const finalDiagnostics = await checkCompilationErrors(tempProjectPath);
                 if (finalDiagnostics.diagnostics && finalDiagnostics.diagnostics.length > 0) {
@@ -297,14 +367,38 @@ 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 });
+
+                // Get final project metrics after generation
+                const finalProjectMetrics = await getProjectMetrics();
+
+                // Send telemetry for generation completion
+                sendTelemetryEvent(extension.ballerinaExtInstance, TM_EVENT_BALLERINA_AI_GENERATION_COMPLETED, CMP_BALLERINA_AI_GENERATION, {
+                    projectId: stateContext.projectId || 'unknown',
+                    messageId: messageId,
+                    modifiedFilesCount: modifiedFiles.length.toString(),
+                    generationStartTime: generationStartTime.toString(),
+                    generationEndTime: generationEndTime.toString(),
+                    durationMs: (generationEndTime - generationStartTime).toString(),
+                    isPlanMode: isPlanModeEnabled.toString(),
+                    approvalMode: stateContext.autoApproveEnabled ? 'auto' : 'manual',
+                    diagnosticChecksCount: diagnosticCheckCount.toString(),
+                    totalCompilationErrorsDuringGeneration: totalCompilationErrorsDuringGeneration.toString(),
+                    finalCompilationErrorsAfterGeneration: (finalDiagnostics.diagnostics?.length || 0).toString(),
+                    inputTokens: inputTokens.toString(),
+                    outputTokens: outputTokens.toString(),
+                    totalTokens: totalTokens.toString(),
+                    outputFileCount: finalProjectMetrics.fileCount.toString(),
+                    outputLineCount: finalProjectMetrics.lineCount.toString(),
+                });
+
                 AIChatStateMachine.sendEvent({
                     type: AIChatMachineEventType.FINISH_EXECUTION,
                 });
                 await langfuseExporter.forceFlush();
                 return tempProjectPath;
             }
         }
-        }
+    }
 
     return tempProjectPath;
 }

workspaces/ballerina/ballerina-extension/src/features/ai/service/feedback/feedback.ts

@@ -0,0 +1,64 @@
+// Copyright (c) 2025, WSO2 LLC. (https://www.wso2.com/) All Rights Reserved.
+
+// WSO2 LLC. licenses this file to you under the Apache License,
+// Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.
+// You may obtain a copy of the License at
+
+// http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import { SubmitFeedbackRequest } from "@wso2/ballerina-core";
+import { fetchWithAuth } from "../connection";
+import { OLD_BACKEND_URL } from "../../utils";
+import { extension } from "../../../../BalExtensionContext";
+import { sendTelemetryEvent, TM_EVENT_BALLERINA_AI_GENERATION_FEEDBACK, CMP_BALLERINA_AI_GENERATION } from "../../../telemetry";
+import { cleanDiagnosticMessages } from "../../../../rpc-managers/ai-panel/utils";
+
+export async function submitFeedback(content: SubmitFeedbackRequest): Promise<boolean> {
+    try {
+        sendTelemetryEvent(
+            extension.ballerinaExtInstance,
+            TM_EVENT_BALLERINA_AI_GENERATION_FEEDBACK,
+            CMP_BALLERINA_AI_GENERATION,
+            {
+                feedbackType: content.positive ? 'positive' : 'negative',
+                hasFeedbackText: content.feedbackText ? 'true' : 'false',
+                feedbackTextLength: content.feedbackText?.length.toString() || '0',
+                hasChatThread: content.messages.length > 0 ? 'true' : 'false',
+                chatThread: JSON.stringify(content.messages),
+            }
+        );
+
+        const payload = {
+            feedback: content.feedbackText,
+            positive: content.positive,
+            messages: content.messages,
+            diagnostics: cleanDiagnosticMessages(content.diagnostics)
+        };
+
+        const response = await fetchWithAuth(`${OLD_BACKEND_URL}/feedback`, {
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/json'
+            },
+            body: JSON.stringify(payload)
+        });
+
+        if (response.ok) {
+            return true;
+        } else {
+            console.error("Failed to submit feedback");
+            return false;
+        }
+    } catch (error) {
+        console.error("Error submitting feedback:", error);
+        return false;
+    }
+}

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

@@ -21,7 +21,7 @@ import { debug } from "../../utils";
 import { window } from "vscode";
 import {
     CMP_EDITOR_SUPPORT, getMessageObject, getTelemetryProperties, sendTelemetryEvent, TM_ERROR_LANG_SERVER,
-    TM_EVENT_EDIT_DIAGRAM, TM_EVENT_EDIT_SOURCE, TM_EVENT_KILL_TERMINAL, TM_FEATURE_USAGE_LANG_SERVER
+    TM_EVENT_EDIT_DIAGRAM, TM_EVENT_EDIT_SOURCE, TM_EVENT_KILL_TERMINAL, TM_FEATURE_USAGE_LANG_SERVER,
 } from ".";
 
 const schedule = require('node-schedule');
@@ -35,25 +35,29 @@ export function activate(ballerinaExtInstance: BallerinaExtension) {
     const langClient = <ExtendedLangClient>ballerinaExtInstance.langClient;
 
     // Start listening telemtry events from language server
-    langClient.onNotification('telemetry/event', (event: LSTelemetryEvent) => {
+    langClient.onNotification('telemetry/event', async (event: LSTelemetryEvent) => {
         let props: { [key: string]: string; };
         switch (event.type) {
             case TM_EVENT_TYPE_ERROR:
                 const errorEvent: LSErrorTelemetryEvent = <LSErrorTelemetryEvent>event;
-                props = getTelemetryProperties(ballerinaExtInstance, event.component, getMessageObject(TM_EVENT_TYPE_ERROR));
+                props = await getTelemetryProperties(ballerinaExtInstance, event.component, getMessageObject(TM_EVENT_TYPE_ERROR));
                 props["ballerina.langserver.error.description"] = errorEvent.message;
                 props["ballerina.langserver.error.stacktrace"] = errorEvent.errorStackTrace;
                 props["ballerina.langserver.error.message"] = errorEvent.errorMessage;
-                reporter.sendTelemetryEvent(TM_ERROR_LANG_SERVER, props);
+
+                // TODO: Enable once when the language server telemerty complete
+                // reporter.sendTelemetryEvent(TM_ERROR_LANG_SERVER, props);
                 break;
             case TM_EVENT_TYPE_FEATURE_USAGE:
                 const usageEvent: LSFeatureUsageTelemetryEvent = <LSFeatureUsageTelemetryEvent>event;
-                props = getTelemetryProperties(ballerinaExtInstance, event.component,
+                props = await getTelemetryProperties(ballerinaExtInstance, event.component,
                     getMessageObject(TM_EVENT_TYPE_FEATURE_USAGE));
                 props["ballerina.langserver.feature.name"] = usageEvent.featureName;
                 props["ballerina.langserver.feature.class"] = usageEvent.featureClass;
                 props["ballerina.langserver.feature.message"] = usageEvent.featureMessage;
-                reporter.sendTelemetryEvent(TM_FEATURE_USAGE_LANG_SERVER, props);
+
+                // TODO: Enable once when the language server telemerty complete
+                // reporter.sendTelemetryEvent(TM_FEATURE_USAGE_LANG_SERVER, props);
                 break;
             default:
                 // Do nothing

workspaces/ballerina/ballerina-extension/src/features/telemetry/common/project-metrics.ts

@@ -0,0 +1,54 @@
+// Copyright (c) 2025, WSO2 LLC. (https://www.wso2.com/) All Rights Reserved.
+
+// WSO2 LLC. licenses this file to you under the Apache License,
+// Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.
+// You may obtain a copy of the License at
+
+// http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import * as vscode from 'vscode';
+import * as fs from 'fs';
+
+export interface ProjectMetrics {
+    fileCount: number;
+    lineCount: number;
+}
+
+export async function getProjectMetrics(): Promise<ProjectMetrics> {
+    const workspaceFolders = vscode.workspace.workspaceFolders;
+
+    if (!workspaceFolders || workspaceFolders.length === 0) {
+        return { fileCount: 0, lineCount: 0 };
+    }
+    const files = await vscode.workspace.findFiles(
+        '**/*.bal',
+        '**/target/**'
+    );
+
+    let totalFileCount = 0;
+    let totalLineCount = 0;
+
+    for (const fileUri of files) {
+        try {
+            totalFileCount++;
+            const fileContent = await fs.promises.readFile(fileUri.fsPath, 'utf8');
+            const lineCount = fileContent.split('\n').length;
+            totalLineCount += lineCount;
+        } catch (error) {
+            console.warn(`Failed to read file ${fileUri.fsPath}:`, error);
+        }
+    }
+
+    return {
+        fileCount: totalFileCount,
+        lineCount: totalLineCount
+    };
+}

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

@@ -45,3 +45,5 @@ export const CMP_CHOREO_AUTHENTICATION = "component.choreo.authentication";
 export const CMP_PERF_ANALYZER = "component.perf.analyzer";
 export const CMP_NOTEBOOK = "component.notebook";
 export const CMP_OPEN_VSCODE_URL = "component.open.vscode.url";
+
+export const CMP_BALLERINA_AI_GENERATION = "ballerina.ai.generation";

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

@@ -108,7 +108,6 @@ export const TM_EVENT_OPEN_VARIABLE_VIEW = "notebook.variable-view.open";
 export const TM_EVENT_UPDATE_VARIABLE_VIEW = "notebook.variable-view.update";
 export const TM_EVENT_START_NOTEBOOK_DEBUG = "notebook.start.debug";
 
-
 // events for open vscode from url
 export const TM_EVENT_OPEN_FILE_URL_START = "vscode.open.file.url.start";
 export const TM_EVENT_OPEN_FILE_CHANGE_PATH = "vscode.open.file.change.path";
@@ -122,3 +121,11 @@ export const TM_EVENT_OPEN_REPO_CHANGE_PATH = "vscode.open.repo.change.path";
 export const TM_EVENT_OPEN_REPO_CANCELED = "vscode.open.repo.canceled";
 export const TM_EVENT_OPEN_REPO_NEW_FOLDER = "vscode.open.exist.repo.new.folder";
 export const TM_EVENT_OPEN_REPO_SAME_FOLDER = "vscode.open.exist.repo.same.folder";
+
+// events for AI features
+export const TM_EVENT_BALLERINA_AI_GENERATION_SUBMITTED = "ballerina.ai.generation.submitted";
+export const TM_EVENT_BALLERINA_AI_GENERATION_COMPLETED = "ballerina.ai.generation.completed";
+export const TM_EVENT_BALLERINA_AI_GENERATION_FAILED = "ballerina.ai.generation.failed";
+export const TM_EVENT_BALLERINA_AI_GENERATION_ABORTED = "ballerina.ai.generation.aborted";
+export const TM_EVENT_BALLERINA_AI_GENERATION_REVERTED = "ballerina.ai.generation.reverted";
+export const  TM_EVENT_BALLERINA_AI_GENERATION_FEEDBACK = "ballerina.ai.generation.feedback";