utils: Changes to enable re running commands with copilot (#2136)

* add re run changes

* add subscription to ActivityAttributes

* changes

* fix build

* oops

* revert

* small changes

* small changes

* requested changes

* bump version

* make copilot sdk optional

* remove

* try removing skipLibCheck

Author: motm32

utils/index.d.ts

@@ -8,6 +8,7 @@
 import type { Environment } from '@azure/ms-rest-azure-env';
 import type { AzExtResourceType, AzureResource, AzureSubscription, ResourceModelBase } from '@microsoft/vscode-azureresources-api';
 import type * as duration from 'dayjs/plugin/duration';
+import type * as vscode from 'vscode';
 import type * as vscodeTypes from 'vscode';
 import { AuthenticationSession, AuthenticationWwwAuthenticateRequest, CancellationToken, CancellationTokenSource, Command, Disposable, Event, ExtensionContext, FileChangeEvent, FileChangeType, FileStat, FileSystemProvider, FileType, InputBoxOptions, LanguageModelToolInvocationOptions, LanguageModelToolInvocationPrepareOptions, LanguageModelToolResult, LogLevel, LogOutputChannel, MarkdownString, MessageItem, MessageOptions, OpenDialogOptions, OutputChannel, PreparedToolInvocation, Progress, ProviderResult, QuickPickItem, TelemetryTrustedValue, TextDocumentShowOptions, ThemeIcon, TreeDataProvider, TreeItem, TreeItemCollapsibleState, TreeView, Uri, QuickPickOptions as VSCodeQuickPickOptions, WorkspaceFolder, WorkspaceFolderPickOptions } from 'vscode';
 import { TargetPopulation } from 'vscode-tas-client';
@@ -707,6 +708,14 @@ export declare function callWithTelemetryAndErrorHandlingSync<T>(callbackId: str
  */
 export declare function callWithMaskHandling<T>(callback: () => Promise<T>, valueToMask: string): Promise<T>;
 
+/**
+ * A wrapper for the VS Code executeCommand command
+ * Used to pass in additional context when executing a command
+ * @param commandId Identifier of the command to execute
+ * @param additionalContext Full context including addtional properties
+ * @param args Parameters passed to the command function
+ */
+export declare function executeCommandWithAddedContext<T>(commandId: string, additionalContext: Partial<IActionContext>, ...args: unknown[]): Thenable<T>;
 /**
  * Add an extension-wide value to mask for all commands
  * This will apply to telemetry and "Report Issue", but _not_ VS Code UI (i.e. the error notification or output channel)
@@ -1413,6 +1422,10 @@ export interface ActivityAttributes {
      * Any Azure resource envelope related to the command or activity being run
      */
     azureResource?: unknown;
+    /**
+     * Optional Azure subscription to be added
+     */
+    subscription?: AzureSubscription;
 
     // For additional one-off properties that could be useful for Copilot
     [key: string]: unknown;
@@ -2383,6 +2396,13 @@ export declare class QuickPickAzureResourceStep extends GenericQuickPickStep<Azu
 export declare function runQuickPickWizard<TPick>(context: PickExperienceContext, wizardOptions?: IWizardOptions<AzureResourceQuickPickWizardContext>, startingNode?: unknown): Promise<TPick>;
 //#endregion
 
+/**
+ * Creates and runs a generic prompt step. Use runQuickPickWizard to run quick pick steps
+ * @param context The action context
+ * @param wizardOptions The options used to construct the wizard
+ */
+export declare function runGenericPromptStep(context: PickExperienceContext, wizardOptions: IWizardOptions<AzureResourceQuickPickWizardContext>): Promise<void>;
+
 /**
  * Registers a namespace for common random utility functions
  */
@@ -2909,6 +2929,22 @@ export declare function runWithInputs<T>(callbackId: string, inputs: (string | R
  */
 export declare function testGlobalSetup(): UIExtensionVariables;
 
+/**
+ * Disposes of copilot session created by `CopilotUserInput`
+ * Should be called after commands using `CopilotUserInput` to prevent any lingering copilot sessions
+ */
+export function disposeCopilotSession(): Promise<void>;
+
+/**
+ * Checks if the context is using `CopilotUserInput`
+ */
+export function isCopilotUserInput(context: IActionContext): boolean;
+
+/**
+ * When setting the ui to `CopilotUserInput`, call this function so that the context can be properly identified
+ * @param context The context to mark as using `CopilotUserInput`
+ */
+export function markAsCopilotUserInput(context: IActionContext, relevantContext?: string, getLoadingView?: () => vscode.WebviewPanel | undefined): void;
 /**
  * Checks if the Copilot CLI is installed, and if not, prompts the user to install it.
  * If the user agrees to install it, this will attempt to install the Copilot CLI automatically.

utils/package-lock.json

@@ -1,7 +1,7 @@
 {
     "name": "@microsoft/vscode-azext-utils",
     "author": "Microsoft Corporation",
-    "version": "4.0.5",
+    "version": "4.0.6",
     "description": "Common UI tools for developing Azure extensions for VS Code",
     "tags": [
         "azure",
@@ -45,7 +45,13 @@
         "vscode-tas-client": "^0.1.84"
     },
     "peerDependencies": {
-        "@azure/ms-rest-azure-env": "^2.0.0"
+        "@azure/ms-rest-azure-env": "^2.0.0",
+        "@github/copilot-sdk": "0.1.32"
+    },
+    "peerDependenciesMeta": {
+        "@github/copilot-sdk": {
+            "optional": true
+        }
     },
     "devDependencies": {
         "@microsoft/vscode-azext-eng": "1.0.0-alpha.12",
@@ -54,6 +60,7 @@
         "@types/semver": "^7.7.1",
         "@types/vscode": "1.105.0",
         "@vscode/test-cli": "*",
-        "@vscode/test-electron": "*"
+        "@vscode/test-electron": "*",
+        "@github/copilot-sdk": "*"
     }
 }

utils/package.json

@@ -3,34 +3,35 @@
 *  Licensed under the MIT License. See License.md in the project root for license information.
 *--------------------------------------------------------------------------------------------*/
 
-import * as vscode from "vscode";
+import type { CopilotClient, CopilotSession } from "@github/copilot-sdk";
+import type * as vscode from "vscode";
 import { InvalidCopilotResponseError } from "../errors";
 
-const languageModelPreference: { vendor: "copilot", family: string }[] = [
-    // not yet seen/available
-    // { vendor: "copilot", family: "gpt-4-turbo-preview" },
-    // seen/available
-    { vendor: "copilot", family: "gpt-4o" },
-    { vendor: "copilot", family: "gpt-4-turbo" },
-    { vendor: "copilot", family: "gpt-4" },
-    { vendor: "copilot", family: "gpt-3.5-turbo" },
-];
-
-async function selectMostPreferredLm(): Promise<vscode.LanguageModelChat | undefined> {
-    const lms = await vscode.lm.selectChatModels({ vendor: "copilot" });
-    const lmsInPreferredOrder = (lms || [])
-        .filter((lm) => languageModelPreference.some((pref) => pref.family === lm.family && pref.vendor === lm.vendor))
-        .sort((a, b) => languageModelPreference.findIndex((pref) => pref.family === a.family && pref.vendor === a.vendor) - languageModelPreference.findIndex((pref) => pref.family === b.family && pref.vendor === b.vendor));
-    return lmsInPreferredOrder?.at(0);
+let client: CopilotClient | undefined;
+let session: CopilotSession | undefined;
+
+async function loadCopilotSdk(): Promise<typeof import("@github/copilot-sdk")> {
+    return await import("@github/copilot-sdk");
 }
 
-export function createPrimaryPromptToGetSingleQuickPickInput(picks: string[], relevantContext?: string): string {
-    return `The User is asking you to choose an item based on the following information:
-        1. Choose from this list of picks ${picks.join(", ")}.
-        2. You must choose one item from the list
-        3. Use information from this context, if available, to make your decision: ${relevantContext ? relevantContext : "No context available"}
-        4. If there is an option to skipForNow, you can choose that option.
-        Respond with a JSON object of the pick you have chosen. Do not respond in a conversational tone, only JSON. `;
+function getCopilotCliPath(): string {
+    return require.resolve(`@github/copilot-${process.platform}-${process.arch}`);
+}
+
+export function createPrimaryPromptToGetSingleQuickPickInput(picks: string[], placeholder?: string): string {
+    return `
+        Task: choose one pick.
+
+        Input:
+        picks: ${JSON.stringify(picks)}
+        placeholder: ${placeholder ?? null}
+
+        Rules:
+        - If the picks represent container image tags, always prefer the newest available image tag.
+
+        Output:
+        Return ONLY the JSON object with fields "label" and "description".
+        `;
 }
 
 export function createPrimaryPromptToGetPickManyQuickPickInput(picks: string[], relevantContext?: string): string {
@@ -63,35 +64,68 @@ export function createPrimaryPromptForWorkspaceFolderPick(folders: readonly vsco
         Respond with the workspace folder you have chosen. Do not respond in a conversational tone. `;
 }
 
-export async function doCopilotInteraction(primaryPrompt: string): Promise<string> {
-    let messages: vscode.LanguageModelChatMessage[] = [];
-    const lm: vscode.LanguageModelChat | undefined = await selectMostPreferredLm();
-
-    messages = [
-        new vscode.LanguageModelChatMessage(vscode.LanguageModelChatMessageRole.User, primaryPrompt)
-    ];
-
-    if (lm !== undefined) {
-        const chatRequestOptions = { justification: `Access to Copilot for the @azure agent.` };
-        const request = await lm.sendRequest(messages, chatRequestOptions);
-        const fragments: string[] = [];
-        try {
-            // Consume the stream and collect fragments
-            for await (const fragment of request.text) {
-                fragments.push(fragment);
-            }
-
-            // Combine fragments into a single string
-            const responseText = fragments.join("");
-            const cleanedResponse = extractJsonString(responseText);
-            return cleanedResponse;
-        } catch {
-            throw new InvalidCopilotResponseError();
-        }
+export async function doGithubCopilotInteraction(primaryPrompt: string, relevantContext?: string): Promise<string> {
+    const session = await getCopilotSession(relevantContext);
+    const response = await session.sendAndWait({
+        prompt: primaryPrompt,
+        mode: "immediate"
+    });
+
+    if (!response || !response.data) {
+        throw new InvalidCopilotResponseError();
+    }
+
+    return response?.data.content;
+}
+
+export async function getCopilotSession(relevantContext?: string): Promise<CopilotSession> {
+    if (session) {
+        return session;
+    }
+
+    const { CopilotClient } = await loadCopilotSdk();
+    client = new CopilotClient({ cliPath: getCopilotCliPath() });
+    session = await client.createSession({
+        onPermissionRequest: () => ({ kind: "approved" })
+    });
+    const activityChildren = extractActivityChildren(relevantContext || '');
+    const subscriptionId = relevantContext ? extractSubscriptionIdFromContext(relevantContext) : undefined;
+
+    await session.sendAndWait({
+        mode: "immediate",
+        prompt: `Picker. Choose one. Never explain your reasoning. Never use markdown. Output: {"label":"X","description":"Y"}
+
+            Rules (priority)
+            1. Match activityChildren "Use X" → pick X
+            2. Match subscriptionId → pick.data /subscriptions/{id}
+            3. Else skipForNow
+
+            Context: ${JSON.stringify({ activityChildren, subscriptionId, relevantContext })}`
+    });
+
+    return session;
+}
+
+export async function disposeCopilotSession(): Promise<void> {
+    if (client) {
+        await client.stop();
+        client = undefined;
+        session = undefined;
     }
-    return "";
 }
 
-function extractJsonString(raw: string): string {
-    return raw.replace(/```json\s*|```/g, '').trim();
+function extractSubscriptionIdFromContext(context: string): string | undefined {
+    const regex = /https:\/\/management\.[^/]+\/subscriptions\/([0-9a-fA-F-]{36})/;
+    const match = context.match(regex);
+    return match ? match[1] : undefined;
+}
+
+function extractActivityChildren(context: string): string | undefined {
+    try {
+        const activityLog = JSON.parse(context) as { children?: unknown };
+        const children = activityLog.children;
+        return JSON.stringify(children);
+    } catch {
+        return undefined;
+    }
 }

utils/src/copilot/copilot.ts

@@ -0,0 +1,11 @@
+/*---------------------------------------------------------------------------------------------
+*  Copyright (c) Microsoft Corporation. All rights reserved.
+*  Licensed under the MIT License. See License.md in the project root for license information.
+*--------------------------------------------------------------------------------------------*/
+import * as vscode from "vscode";
+import * as types from '../index';
+
+export function executeCommandWithAddedContext<T = unknown>(commandId: string, additionalContext: Partial<types.IActionContext>, ...args: unknown[]): Thenable<T> {
+    const metadata = { __injectedContext: additionalContext };
+    return vscode.commands.executeCommand(commandId, ...args, metadata);
+}

utils/src/executeCommandWithAddedContext.ts

@@ -9,6 +9,7 @@ export { createAzExtLogOutputChannel, createAzExtOutputChannel } from './AzExtOu
 export * from './AzExtTreeFileSystem';
 export * from './callWithTelemetryAndErrorHandling';
 export { activityErrorContext, activityFailContext, activityFailIcon, activityInfoContext, activityInfoIcon, activityProgressContext, activityProgressIcon, activitySuccessContext, activitySuccessIcon } from './constants';
+export * from './copilot/copilot';
 export * from './copilot/installCopilotCli';
 export * from './createApiProvider';
 export { createExperimentationService } from './createExperimentationService';
@@ -18,6 +19,7 @@ export * from './dev/TestOutputChannel';
 export * from './dev/TestUserInput';
 export * from './DialogResponses';
 export * from './errors';
+export * from './executeCommandWithAddedContext';
 export * from './extensionUserAgent';
 export { registerUIExtensionVariables } from './extensionVariables';
 export { addExtensionValueToMask, callWithMaskHandling, maskUserInfo, maskValue } from './masking';
@@ -33,6 +35,7 @@ export * from './pickTreeItem/GenericQuickPickStep';
 export * from './pickTreeItem/quickPickAzureResource/QuickPickAzureResourceStep';
 export * from './pickTreeItem/quickPickAzureResource/QuickPickAzureSubscriptionStep';
 export * from './pickTreeItem/quickPickAzureResource/QuickPickGroupStep';
+export * from './pickTreeItem/runGenericPromptStep';
 export * from './pickTreeItem/runQuickPickWizard';
 export * from './registerCommand';
 export * from './registerCommandWithTreeNodeUnwrapping';
@@ -55,6 +58,7 @@ export * from './utils/activityUtils';
 export * from './utils/AzExtFsExtra';
 export * from './utils/AzureResourceIdTelemetry';
 export * from './utils/contextUtils';
+export * from './utils/copilotUtils';
 export * from './utils/credentialUtils';
 export * from './utils/dateTimeUtils';
 export * from './utils/findFreePort';

utils/src/index.ts

@@ -0,0 +1,17 @@
+/*---------------------------------------------------------------------------------------------
+*  Copyright (c) Microsoft Corporation. All rights reserved.
+*  Licensed under the MIT License. See License.md in the project root for license information.
+*--------------------------------------------------------------------------------------------*/
+
+import * as types from '../../index';
+import { AzureWizard } from '../wizard/AzureWizard';
+
+export async function runGenericPromptStep(context: types.PickExperienceContext, wizardOptions?: types.IWizardOptions<types.AzureResourceQuickPickWizardContext>): Promise<void> {
+    const wizard = new AzureWizard(context, {
+        ...wizardOptions,
+        hideStepCount: wizardOptions?.hideStepCount ?? true,
+        showLoadingPrompt: wizardOptions?.showLoadingPrompt ?? true
+    });
+
+    await wizard.prompt();
+}

utils/src/pickTreeItem/runGenericPromptStep.ts

@@ -42,7 +42,20 @@ export function registerCommand(commandId: string, callback: (context: types.IAc
         return await callWithTelemetryAndErrorHandling(
             telemetryId || commandId,
             async (context: types.IActionContext) => {
+                let injectedContext: Partial<types.IActionContext> | undefined;
                 if (args.length > 0) {
+                    const metadata = args.find(
+                        (arg) => arg !== null && typeof arg === "object" && Object.prototype.hasOwnProperty.call(arg, "__injectedContext")
+                    );
+
+                    if (metadata) {
+                        const candidate = (metadata as Record<string, unknown>).__injectedContext;
+                        if (candidate && typeof candidate === "object") {
+                            injectedContext = candidate as Partial<types.IActionContext>;
+                        }
+                        args.splice(args.indexOf(metadata), 1); // remove only the metadata
+                    }
+
                     try {
                         await setTelemetryProperties(context, args);
                     } catch (e: unknown) {
@@ -52,8 +65,8 @@ export function registerCommand(commandId: string, callback: (context: types.IAc
                         context.telemetry.properties.telemetryError = error.message;
                     }
                 }
-
-                return callback(context, ...args);
+                const finalContext = injectedContext ? { ...context, ...injectedContext } : context;
+                return callback(finalContext, ...args);
             }
         );
     }));