Skip to content
36 changes: 36 additions & 0 deletions utils/index.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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
*/
Expand Down Expand Up @@ -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.
Expand Down
2,327 changes: 1,493 additions & 834 deletions utils/package-lock.json

Large diffs are not rendered by default.

13 changes: 10 additions & 3 deletions utils/package.json
Original file line number Diff line number Diff line change
@@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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": "*"
}
}
138 changes: 86 additions & 52 deletions utils/src/copilot/copilot.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

relevantContext is ignored if session already exists due to the logic on line 68. Might want to redesign this so that session isn't a global variable.

Copy link
Copy Markdown
Contributor Author

@motm32 motm32 Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is by design we only want to pass in the relevant context when we create the session. In the subsequent calls it will be using that context and the prompt is changed through createPrimaryPromptToGetSingleQuickPickInput so that the proper placeholder and picks get sent each time. If there is no session found then we need to send the relevantContext when we create the session.

I designed it this way as it is much faster to not have create a new session on each quick pick call and send in all the relevant context everytime.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of storing the Copilot session as a module-level global, it might be best to move the session ownership into the CopilotUserInput class itself.

Each CopilotUserInput instance could lazily create its own CopilotSession on the first interaction and reuse it for subsequent wizard steps within the same command so we keep the performance benefit of not recreating sessions for each step. But since each command already gets its own CopilotUserInput via context.ui, the session is naturally scoped to that command's lifecycle. This removes the need for callers to remember to call disposeCopilotSession(), and makes concurrent commands safe since they'd each have their own session. Also, the doGithubCopilotInteraction function would take a session as a parameter instead of managing the global itself.

const response = await session.sendAndWait({
prompt: primaryPrompt,
mode: "immediate"
});

if (!response || !response.data) {
throw new InvalidCopilotResponseError();
}

return response?.data.content;
}
Comment on lines +67 to +79
Copy link

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

doGithubCopilotInteraction returns '' when no content is returned. Call sites then JSON.parse the response and will throw. Consider throwing an InvalidCopilotResponseError (or similar) when response?.data.content is falsy, and strip common wrappers (e.g., ```json fences) like the previous implementation did to make parsing more reliable.

Copilot uses AI. Check for mistakes.

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;
Comment on lines +81 to +92
Copy link

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

getCopilotSession caches a single global session and returns it even when relevantContext changes. This can cause subsequent prompts to be evaluated using stale/incorrect context (and can leak context between commands). Consider scoping the session to a command invocation, or tracking the last relevantContext and recreating/updating the session when it differs.

Copilot uses AI. Check for mistakes.

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 {
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is going to be too strict for any other cloud (government, China, etc.)

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;
}

It's a little looser matching, but I don't think we need to be too strict since we kind of know what we're parsing.

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;
}
}
11 changes: 11 additions & 0 deletions utils/src/executeCommandWithAddedContext.ts
Original file line number Diff line number Diff line change
@@ -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);
}
Comment on lines +8 to +11
Copy link

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There’s no automated coverage for the new executeCommandWithAddedContext + registerCommand injection path. Given this changes how command args/context are interpreted, add unit tests to ensure: metadata is stripped from args, context is merged as intended, and normal command calls (including object args) aren’t mis-identified as metadata.

Copilot uses AI. Check for mistakes.
4 changes: 4 additions & 0 deletions utils/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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';
Expand All @@ -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';
Expand All @@ -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';
Expand Down
17 changes: 17 additions & 0 deletions utils/src/pickTreeItem/runGenericPromptStep.ts
Original file line number Diff line number Diff line change
@@ -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();
}
17 changes: 15 additions & 2 deletions utils/src/registerCommand.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +47 to +56
Copy link

Copilot AI Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The injected-context detection uses "__injectedContext" in arg, which will match inherited properties and can be spoofed by any caller invoking the command with an object arg. Consider only accepting metadata when it is the last argument and Object.prototype.hasOwnProperty.call(arg, "__injectedContext") is true, and validate that __injectedContext is a plain object before merging. Also, the cast to types.IActionContext is too strong here—this is a Partial<IActionContext>.

Copilot uses AI. Check for mistakes.
}

try {
await setTelemetryProperties(context, args);
} catch (e: unknown) {
Expand All @@ -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);
}
);
}));
Expand Down
Loading
Loading