More tests for Copilot CLI and refactor permission handling

src/extension/agents/copilotcli/node/copilotCli.ts

@@ -12,13 +12,19 @@ import { ILogService } from '../../../../platform/log/common/logService';
 import { IWorkspaceService } from '../../../../platform/workspace/common/workspaceService';
 import { createServiceIdentifier } from '../../../../util/common/services';
 import { Lazy } from '../../../../util/vs/base/common/lazy';
-import { Disposable, IDisposable, toDisposable } from '../../../../util/vs/base/common/lifecycle';
+import { IDisposable, toDisposable } from '../../../../util/vs/base/common/lifecycle';
 import { getCopilotLogger } from './logger';
 import { ensureNodePtyShim } from './nodePtyShim';
+import { PermissionRequest } from './permissionHelpers';
 
 const COPILOT_CLI_MODEL_MEMENTO_KEY = 'github.copilot.cli.sessionModel';
 const DEFAULT_CLI_MODEL = 'claude-sonnet-4';
 
+export interface CopilotCLISessionOptions {
+	addPermissionHandler(handler: SessionOptions['requestPermission']): IDisposable;
+	toSessionOptions(): SessionOptions;
+}
+
 export interface ICopilotCLIModels {
 	_serviceBrand: undefined;
 	toModelProvider(modelId: string): string;
@@ -104,9 +110,8 @@ export class CopilotCLISDK implements ICopilotCLISDK {
 export interface ICopilotCLISessionOptionsService {
 	readonly _serviceBrand: undefined;
 	createOptions(
-		options: SessionOptions,
-		permissionHandler: CopilotCLIPermissionsHandler
-	): Promise<SessionOptions>;
+		options: SessionOptions
+	): Promise<CopilotCLISessionOptions>;
 }
 export const ICopilotCLISessionOptionsService = createServiceIdentifier<ICopilotCLISessionOptionsService>('ICopilotCLISessionOptionsService');
 
@@ -118,17 +123,28 @@ export class CopilotCLISessionOptionsService implements ICopilotCLISessionOption
 		@ILogService private readonly logService: ILogService,
 	) { }
 
-	public async createOptions(options: SessionOptions, permissionHandler: CopilotCLIPermissionsHandler) {
+	public async createOptions(options: SessionOptions) {
 		const copilotToken = await this._authenticationService.getAnyGitHubSession();
 		const workingDirectory = options.workingDirectory ?? await this.getWorkspaceFolderPath();
+		const logger = this.logService;
+		const requestPermissionRejected = async (permission: PermissionRequest): ReturnType<NonNullable<SessionOptions['requestPermission']>> => {
+			logger.info(`[CopilotCLISessionOptionsService] Permission request denied for permission as no handler was set: ${permission.kind}`);
+			return {
+				kind: "denied-interactively-by-user"
+			};
+		};
+		const permissionHandler: Required<Pick<SessionOptions, 'requestPermission'>> = {
+			requestPermission: requestPermissionRejected
+		};
+
 		const allOptions: SessionOptions = {
 			env: {
 				...process.env,
 				COPILOTCLI_DISABLE_NONESSENTIAL_TRAFFIC: '1'
 			},
 			logger: getCopilotLogger(this.logService),
-			requestPermission: async (permissionRequest) => {
-				return await permissionHandler.getPermissions(permissionRequest);
+			requestPermission: async (request: PermissionRequest) => {
+				return await permissionHandler.requestPermission(request);
 			},
 			authInfo: {
 				type: 'token',
@@ -141,7 +157,18 @@ export class CopilotCLISessionOptionsService implements ICopilotCLISessionOption
 		if (workingDirectory) {
 			allOptions.workingDirectory = workingDirectory;
 		}
-		return allOptions;
+
+		return {
+			addPermissionHandler: (handler: NonNullable<SessionOptions['requestPermission']>) => {
+				permissionHandler.requestPermission = handler;
+				return toDisposable(() => {
+					if (permissionHandler.requestPermission === handler) {
+						permissionHandler.requestPermission = requestPermissionRejected;
+					}
+				});
+			},
+			toSessionOptions: () => allOptions
+		} satisfies CopilotCLISessionOptions;
 	}
 	private async getWorkspaceFolderPath() {
 		if (this.workspaceService.getWorkspaceFolders().length === 0) {
@@ -154,32 +181,3 @@ export class CopilotCLISessionOptionsService implements ICopilotCLISessionOption
 		return folder?.uri?.fsPath;
 	}
 }
-
-
-/**
- * Perhaps temporary interface to handle permission requests from the Copilot CLI SDK
- * Perhaps because the SDK needs a better way to handle this in long term per session.
- */
-export interface ICopilotCLIPermissions {
-	onDidRequestPermissions(handler: SessionOptions['requestPermission']): IDisposable;
-}
-
-export class CopilotCLIPermissionsHandler extends Disposable implements ICopilotCLIPermissions {
-	private _handler: SessionOptions['requestPermission'] | undefined;
-
-	public onDidRequestPermissions(handler: SessionOptions['requestPermission']): IDisposable {
-		this._handler = handler;
-		return this._register(toDisposable(() => {
-			this._handler = undefined;
-		}));
-	}
-
-	public async getPermissions(permission: Parameters<NonNullable<SessionOptions['requestPermission']>>[0]): Promise<ReturnType<NonNullable<SessionOptions['requestPermission']>>> {
-		if (!this._handler) {
-			return {
-				kind: "denied-interactively-by-user"
-			};
-		}
-		return await this._handler(permission);
-	}
-}

src/extension/agents/copilotcli/node/copilotcliSession.ts

@@ -3,39 +3,45 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import type { Attachment, Session, SessionOptions } from '@github/copilot/sdk';
+import type { Attachment, Session } from '@github/copilot/sdk';
 import type * as vscode from 'vscode';
 import { ILogService } from '../../../../platform/log/common/logService';
 import { IWorkspaceService } from '../../../../platform/workspace/common/workspaceService';
 import { CancellationToken } from '../../../../util/vs/base/common/cancellation';
+import { Emitter, Event } from '../../../../util/vs/base/common/event';
 import { DisposableStore, IDisposable, toDisposable } from '../../../../util/vs/base/common/lifecycle';
 import { ResourceMap } from '../../../../util/vs/base/common/map';
 import { extUriBiasedIgnorePathCase } from '../../../../util/vs/base/common/resources';
-import { ChatRequestTurn2, ChatResponseThinkingProgressPart, ChatResponseTurn2, ChatSessionStatus, EventEmitter, LanguageModelTextPart, Uri } from '../../../../vscodeTypes';
-import { IToolsService } from '../../../tools/common/toolsService';
+import { ChatRequestTurn2, ChatResponseThinkingProgressPart, ChatResponseTurn2, ChatSessionStatus, EventEmitter, Uri } from '../../../../vscodeTypes';
 import { ExternalEditTracker } from '../../common/externalEditTracker';
-import { CopilotCLIPermissionsHandler, ICopilotCLISessionOptionsService } from './copilotCli';
+import { CopilotCLISessionOptions, ICopilotCLISessionOptionsService } from './copilotCli';
 import { buildChatHistoryFromEvents, getAffectedUrisForEditTool, isCopilotCliEditToolCall, processToolExecutionComplete, processToolExecutionStart } from './copilotcliToolInvocationFormatter';
-import { getConfirmationToolParams, PermissionRequest } from './permissionHelpers';
+import { PermissionRequest } from './permissionHelpers';
+
+type PermissionHandler = (
+	permissionRequest: PermissionRequest,
+	token: CancellationToken,
+) => Promise<boolean>;
 
 export interface ICopilotCLISession extends IDisposable {
 	readonly sessionId: string;
 	readonly status: vscode.ChatSessionStatus | undefined;
 	readonly onDidChangeStatus: vscode.Event<vscode.ChatSessionStatus | undefined>;
+	readonly permissionRequested?: PermissionRequest;
+	readonly onPermissionRequested: vscode.Event<PermissionRequest>;
 
+	attachPermissionHandler(handler: PermissionHandler): IDisposable;
+	attachStream(stream: vscode.ChatResponseStream): IDisposable;
 	handleRequest(
 		prompt: string,
 		attachments: Attachment[],
 		modelId: string | undefined,
-		stream: vscode.ChatResponseStream,
-		toolInvocationToken: vscode.ChatParticipantToolToken,
 		token: vscode.CancellationToken
 	): Promise<void>;
-
 	addUserMessage(content: string): void;
 	addUserAssistantMessage(content: string): void;
 	getSelectedModelId(): Promise<string | undefined>;
-	getChatHistory(): Promise<(ChatRequestTurn2 | ChatResponseTurn2)[]>;
+	getChatHistory(): (ChatRequestTurn2 | ChatResponseTurn2)[];
 }
 
 export class CopilotCLISession extends DisposableStore implements ICopilotCLISession {
@@ -49,25 +55,49 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
 
 	public readonly onDidChangeStatus = this._statusChange.event;
 
+	private _permissionRequested?: PermissionRequest;
+	public get permissionRequested(): PermissionRequest | undefined {
+		return this._permissionRequested;
+	}
+	private readonly _onPermissionRequested = this.add(new EventEmitter<PermissionRequest>());
+	public readonly onPermissionRequested = this._onPermissionRequested.event;
+	private _permissionHandler?: PermissionHandler;
+	private readonly _permissionHandlerSet = this.add(new Emitter<void>());
+	private _stream?: vscode.ChatResponseStream;
 	constructor(
+		private readonly _options: CopilotCLISessionOptions,
 		private readonly _sdkSession: Session,
-		private readonly _options: SessionOptions,
-		private readonly _permissionHandler: CopilotCLIPermissionsHandler,
 		@ILogService private readonly logService: ILogService,
 		@IWorkspaceService private readonly workspaceService: IWorkspaceService,
-		@IToolsService private readonly toolsService: IToolsService,
 		@ICopilotCLISessionOptionsService private readonly cliSessionOptions: ICopilotCLISessionOptionsService,
 	) {
 		super();
 		this.sessionId = _sdkSession.sessionId;
 	}
 
+	attachStream(stream: vscode.ChatResponseStream): IDisposable {
+		this._stream = stream;
+		return toDisposable(() => {
+			if (this._stream === stream) {
+				this._stream = undefined;
+			}
+		});
+	}
+
+	attachPermissionHandler(handler: PermissionHandler): IDisposable {
+		this._permissionHandler = handler;
+		this._permissionHandlerSet.fire();
+		return toDisposable(() => {
+			if (this._permissionHandler === handler) {
+				this._permissionHandler = undefined;
+			}
+		});
+	}
+
 	public async handleRequest(
 		prompt: string,
 		attachments: Attachment[],
 		modelId: string | undefined,
-		stream: vscode.ChatResponseStream,
-		toolInvocationToken: vscode.ChatParticipantToolToken,
 		token: vscode.CancellationToken
 	): Promise<void> {
 		if (this.isDisposed) {
@@ -88,26 +118,25 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
 		const editToolIds = new Set<string>();
 		const editTracker = new ExternalEditTracker();
 		const editFilesAndToolCallIds = new ResourceMap<string[]>();
-		disposables.add(this._permissionHandler.onDidRequestPermissions(async (permissionRequest) => {
-			return await this.requestPermission(permissionRequest, stream, editTracker,
+		disposables.add(this._options.addPermissionHandler(async (permissionRequest) => {
+			// Need better API from SDK to correlate file edits in permission requests to tool invocations.
+			return await this.requestPermission(permissionRequest, editTracker,
 				(file: Uri) => {
 					const ids = editFilesAndToolCallIds.get(file);
 					return ids?.shift();
 				},
-				toolInvocationToken,
-				this._options.workingDirectory
+				this._options.toSessionOptions().workingDirectory,
+				token
 			);
 		}));
 
 		try {
-			const [currentModel,
-				sessionOptions
-			] = await Promise.all([
+			const [currentModel, authInfo] = await Promise.all([
 				modelId ? this._sdkSession.getSelectedModel() : undefined,
-				this.cliSessionOptions.createOptions(this._options, this._permissionHandler)
+				this.cliSessionOptions.createOptions({}).then(opts => opts.toSessionOptions().authInfo)
 			]);
-			if (sessionOptions.authInfo) {
-				this._sdkSession.setAuthInfo(sessionOptions.authInfo);
+			if (authInfo) {
+				this._sdkSession.setAuthInfo(authInfo);
 			}
 			if (modelId && modelId !== currentModel) {
 				await this._sdkSession.setSelectedModel(modelId);
@@ -116,7 +145,7 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
 			disposables.add(toDisposable(this._sdkSession.on('*', (event) => this.logService.trace(`[CopilotCLISession]CopilotCLI Event: ${JSON.stringify(event, null, 2)}`))));
 			disposables.add(toDisposable(this._sdkSession.on('assistant.message', (event) => {
 				if (typeof event.data.content === 'string' && event.data.content.length) {
-					stream.markdown(event.data.content);
+					this._stream?.markdown(event.data.content);
 				}
 			})));
 			disposables.add(toDisposable(this._sdkSession.on('tool.execution_start', (event) => {
@@ -136,7 +165,7 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
 				} else {
 					const responsePart = processToolExecutionStart(event, this._pendingToolInvocations);
 					if (responsePart instanceof ChatResponseThinkingProgressPart) {
-						stream.push(responsePart);
+						this._stream?.push(responsePart);
 					}
 				}
 				this.logService.trace(`[CopilotCLISession] Start Tool ${event.data.toolName || '<unknown>'}`);
@@ -151,7 +180,7 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
 
 				const responsePart = processToolExecutionComplete(event, this._pendingToolInvocations);
 				if (responsePart && !(responsePart instanceof ChatResponseThinkingProgressPart)) {
-					stream.push(responsePart);
+					this._stream?.push(responsePart);
 				}
 
 				const toolName = toolNames.get(event.data.toolCallId) || '<unknown>';
@@ -163,7 +192,7 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
 			})));
 			disposables.add(toDisposable(this._sdkSession.on('session.error', (event) => {
 				this.logService.error(`[CopilotCLISession]CopilotCLI error: (${event.data.errorType}), ${event.data.message}`);
-				stream.markdown(`\n\n❌ Error: (${event.data.errorType}) ${event.data.message}`);
+				this._stream?.markdown(`\n\n❌ Error: (${event.data.errorType}) ${event.data.message}`);
 			})));
 
 			await this._sdkSession.send({ prompt, attachments, abortController });
@@ -175,7 +204,7 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
 			this._status = ChatSessionStatus.Failed;
 			this._statusChange.fire(this._status);
 			this.logService.error(`[CopilotCLISession] Invoking session (error) ${this.sessionId}`, error);
-			stream.markdown(`\n\n❌ Error: ${error instanceof Error ? error.message : String(error)}`);
+			this._stream?.markdown(`\n\n❌ Error: ${error instanceof Error ? error.message : String(error)}`);
 		} finally {
 			disposables.dispose();
 		}
@@ -196,18 +225,17 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
 		return this._sdkSession.getSelectedModel();
 	}
 
-	public async getChatHistory(): Promise<(ChatRequestTurn2 | ChatResponseTurn2)[]> {
-		const events = await this._sdkSession.getEvents();
+	public getChatHistory(): (ChatRequestTurn2 | ChatResponseTurn2)[] {
+		const events = this._sdkSession.getEvents();
 		return buildChatHistoryFromEvents(events);
 	}
 
 	private async requestPermission(
 		permissionRequest: PermissionRequest,
-		stream: vscode.ChatResponseStream,
 		editTracker: ExternalEditTracker,
 		getEditKeyForFile: (file: Uri) => string | undefined,
-		toolInvocationToken: vscode.ChatParticipantToolToken,
-		workingDirectory?: string
+		workingDirectory: string | undefined,
+		token: vscode.CancellationToken
 	): Promise<{ kind: 'approved' } | { kind: 'denied-interactively-by-user' }> {
 		if (permissionRequest.kind === 'read') {
 			// If user is reading a file in the working directory or workspace, auto-approve
@@ -239,26 +267,40 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
 		}
 
 		try {
-			const { tool, input } = getConfirmationToolParams(permissionRequest);
-			const result = await this.toolsService.invokeTool(tool,
-				{ input, toolInvocationToken },
-				CancellationToken.None);
+			const permissionHandler = await this.waitForPermissionHandler(permissionRequest);
+			if (!permissionHandler) {
+				this.logService.warn(`[CopilotCLISession] No permission handler registered, denying request for ${permissionRequest.kind} permission.`);
+				return { kind: 'denied-interactively-by-user' };
+			}
 
-			const firstResultPart = result.content.at(0);
-			if (firstResultPart instanceof LanguageModelTextPart && firstResultPart.value === 'yes') {
+			if (await permissionHandler(permissionRequest, token)) {
 				// If we're editing a file, start tracking the edit & wait for core to acknowledge it.
 				const editFile = permissionRequest.kind === 'write' ? Uri.file(permissionRequest.fileName) : undefined;
 				const editKey = editFile ? getEditKeyForFile(editFile) : undefined;
-				if (editFile && editKey) {
+				if (editFile && editKey && this._stream) {
 					this.logService.trace(`[CopilotCLISession] Starting to track edit for toolCallId ${editKey} & file ${editFile.fsPath}`);
-					await editTracker.trackEdit(editKey, [editFile], stream);
+					await editTracker.trackEdit(editKey, [editFile], this._stream);
 				}
 				return { kind: 'approved' };
 			}
 		} catch (error) {
 			this.logService.error(`[CopilotCLISession] Permission request error: ${error}`);
+		} finally {
+			this._permissionRequested = undefined;
 		}
 
 		return { kind: 'denied-interactively-by-user' };
 	}
+
+	private async waitForPermissionHandler(permissionRequest: PermissionRequest): Promise<PermissionHandler | undefined> {
+		if (!this._permissionHandler) {
+			this._permissionRequested = permissionRequest;
+			this._onPermissionRequested.fire(permissionRequest);
+			const disposables = this.add(new DisposableStore());
+			await Event.toPromise(this._permissionHandlerSet.event, disposables);
+			disposables.dispose();
+			this._permissionRequested = undefined;
+		}
+		return this._permissionHandler;
+	}
 }