Skip to content

Commit 5c67247

Browse files
DonJayamanneDonJayamanne
authored
More tests for Copilot CLI and refactor permission handling (#1871)
* More Copilot CLI tests and refactor permissions * Updates * Address review comments * Updates * Address review comments * Address review comments
1 parent 9564ab1 commit 5c67247

File tree

9 files changed

+761
-279
lines changed

9 files changed

+761
-279
lines changed

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

Lines changed: 35 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,19 @@ import { ILogService } from '../../../../platform/log/common/logService';
1212
import { IWorkspaceService } from '../../../../platform/workspace/common/workspaceService';
1313
import { createServiceIdentifier } from '../../../../util/common/services';
1414
import { Lazy } from '../../../../util/vs/base/common/lazy';
15-
import { Disposable, IDisposable, toDisposable } from '../../../../util/vs/base/common/lifecycle';
15+
import { IDisposable, toDisposable } from '../../../../util/vs/base/common/lifecycle';
1616
import { getCopilotLogger } from './logger';
1717
import { ensureNodePtyShim } from './nodePtyShim';
18+
import { PermissionRequest } from './permissionHelpers';
1819

1920
const COPILOT_CLI_MODEL_MEMENTO_KEY = 'github.copilot.cli.sessionModel';
2021
const DEFAULT_CLI_MODEL = 'claude-sonnet-4';
2122

23+
export interface CopilotCLISessionOptions {
24+
addPermissionHandler(handler: SessionOptions['requestPermission']): IDisposable;
25+
toSessionOptions(): SessionOptions;
26+
}
27+
2228
export interface ICopilotCLIModels {
2329
_serviceBrand: undefined;
2430
toModelProvider(modelId: string): string;
@@ -104,9 +110,8 @@ export class CopilotCLISDK implements ICopilotCLISDK {
104110
export interface ICopilotCLISessionOptionsService {
105111
readonly _serviceBrand: undefined;
106112
createOptions(
107-
options: SessionOptions,
108-
permissionHandler: CopilotCLIPermissionsHandler
109-
): Promise<SessionOptions>;
113+
options: SessionOptions
114+
): Promise<CopilotCLISessionOptions>;
110115
}
111116
export const ICopilotCLISessionOptionsService = createServiceIdentifier<ICopilotCLISessionOptionsService>('ICopilotCLISessionOptionsService');
112117

@@ -118,17 +123,28 @@ export class CopilotCLISessionOptionsService implements ICopilotCLISessionOption
118123
@ILogService private readonly logService: ILogService,
119124
) { }
120125

121-
public async createOptions(options: SessionOptions, permissionHandler: CopilotCLIPermissionsHandler) {
126+
public async createOptions(options: SessionOptions) {
122127
const copilotToken = await this._authenticationService.getAnyGitHubSession();
123128
const workingDirectory = options.workingDirectory ?? await this.getWorkspaceFolderPath();
129+
const logger = this.logService;
130+
const requestPermissionRejected = async (permission: PermissionRequest): ReturnType<NonNullable<SessionOptions['requestPermission']>> => {
131+
logger.info(`[CopilotCLISessionOptionsService] Permission request denied for permission as no handler was set: ${permission.kind}`);
132+
return {
133+
kind: "denied-interactively-by-user"
134+
};
135+
};
136+
const permissionHandler: Required<Pick<SessionOptions, 'requestPermission'>> = {
137+
requestPermission: requestPermissionRejected
138+
};
139+
124140
const allOptions: SessionOptions = {
125141
env: {
126142
...process.env,
127143
COPILOTCLI_DISABLE_NONESSENTIAL_TRAFFIC: '1'
128144
},
129145
logger: getCopilotLogger(this.logService),
130-
requestPermission: async (permissionRequest) => {
131-
return await permissionHandler.getPermissions(permissionRequest);
146+
requestPermission: async (request: PermissionRequest) => {
147+
return await permissionHandler.requestPermission(request);
132148
},
133149
authInfo: {
134150
type: 'token',
@@ -141,7 +157,18 @@ export class CopilotCLISessionOptionsService implements ICopilotCLISessionOption
141157
if (workingDirectory) {
142158
allOptions.workingDirectory = workingDirectory;
143159
}
144-
return allOptions;
160+
161+
return {
162+
addPermissionHandler: (handler: NonNullable<SessionOptions['requestPermission']>) => {
163+
permissionHandler.requestPermission = handler;
164+
return toDisposable(() => {
165+
if (permissionHandler.requestPermission === handler) {
166+
permissionHandler.requestPermission = requestPermissionRejected;
167+
}
168+
});
169+
},
170+
toSessionOptions: () => allOptions
171+
} satisfies CopilotCLISessionOptions;
145172
}
146173
private async getWorkspaceFolderPath() {
147174
if (this.workspaceService.getWorkspaceFolders().length === 0) {
@@ -154,32 +181,3 @@ export class CopilotCLISessionOptionsService implements ICopilotCLISessionOption
154181
return folder?.uri?.fsPath;
155182
}
156183
}
157-
158-
159-
/**
160-
* Perhaps temporary interface to handle permission requests from the Copilot CLI SDK
161-
* Perhaps because the SDK needs a better way to handle this in long term per session.
162-
*/
163-
export interface ICopilotCLIPermissions {
164-
onDidRequestPermissions(handler: SessionOptions['requestPermission']): IDisposable;
165-
}
166-
167-
export class CopilotCLIPermissionsHandler extends Disposable implements ICopilotCLIPermissions {
168-
private _handler: SessionOptions['requestPermission'] | undefined;
169-
170-
public onDidRequestPermissions(handler: SessionOptions['requestPermission']): IDisposable {
171-
this._handler = handler;
172-
return this._register(toDisposable(() => {
173-
this._handler = undefined;
174-
}));
175-
}
176-
177-
public async getPermissions(permission: Parameters<NonNullable<SessionOptions['requestPermission']>>[0]): Promise<ReturnType<NonNullable<SessionOptions['requestPermission']>>> {
178-
if (!this._handler) {
179-
return {
180-
kind: "denied-interactively-by-user"
181-
};
182-
}
183-
return await this._handler(permission);
184-
}
185-
}

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

Lines changed: 84 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -3,39 +3,45 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6-
import type { Attachment, Session, SessionOptions } from '@github/copilot/sdk';
6+
import type { Attachment, Session } from '@github/copilot/sdk';
77
import type * as vscode from 'vscode';
88
import { ILogService } from '../../../../platform/log/common/logService';
99
import { IWorkspaceService } from '../../../../platform/workspace/common/workspaceService';
1010
import { CancellationToken } from '../../../../util/vs/base/common/cancellation';
11+
import { Emitter, Event } from '../../../../util/vs/base/common/event';
1112
import { DisposableStore, IDisposable, toDisposable } from '../../../../util/vs/base/common/lifecycle';
1213
import { ResourceMap } from '../../../../util/vs/base/common/map';
1314
import { extUriBiasedIgnorePathCase } from '../../../../util/vs/base/common/resources';
14-
import { ChatRequestTurn2, ChatResponseThinkingProgressPart, ChatResponseTurn2, ChatSessionStatus, EventEmitter, LanguageModelTextPart, Uri } from '../../../../vscodeTypes';
15-
import { IToolsService } from '../../../tools/common/toolsService';
15+
import { ChatRequestTurn2, ChatResponseThinkingProgressPart, ChatResponseTurn2, ChatSessionStatus, EventEmitter, Uri } from '../../../../vscodeTypes';
1616
import { ExternalEditTracker } from '../../common/externalEditTracker';
17-
import { CopilotCLIPermissionsHandler, ICopilotCLISessionOptionsService } from './copilotCli';
17+
import { CopilotCLISessionOptions, ICopilotCLISessionOptionsService } from './copilotCli';
1818
import { buildChatHistoryFromEvents, getAffectedUrisForEditTool, isCopilotCliEditToolCall, processToolExecutionComplete, processToolExecutionStart } from './copilotcliToolInvocationFormatter';
19-
import { getConfirmationToolParams, PermissionRequest } from './permissionHelpers';
19+
import { PermissionRequest } from './permissionHelpers';
20+
21+
type PermissionHandler = (
22+
permissionRequest: PermissionRequest,
23+
token: CancellationToken,
24+
) => Promise<boolean>;
2025

2126
export interface ICopilotCLISession extends IDisposable {
2227
readonly sessionId: string;
2328
readonly status: vscode.ChatSessionStatus | undefined;
2429
readonly onDidChangeStatus: vscode.Event<vscode.ChatSessionStatus | undefined>;
30+
readonly permissionRequested?: PermissionRequest;
31+
readonly onPermissionRequested: vscode.Event<PermissionRequest>;
2532

33+
attachPermissionHandler(handler: PermissionHandler): IDisposable;
34+
attachStream(stream: vscode.ChatResponseStream): IDisposable;
2635
handleRequest(
2736
prompt: string,
2837
attachments: Attachment[],
2938
modelId: string | undefined,
30-
stream: vscode.ChatResponseStream,
31-
toolInvocationToken: vscode.ChatParticipantToolToken,
3239
token: vscode.CancellationToken
3340
): Promise<void>;
34-
3541
addUserMessage(content: string): void;
3642
addUserAssistantMessage(content: string): void;
3743
getSelectedModelId(): Promise<string | undefined>;
38-
getChatHistory(): Promise<(ChatRequestTurn2 | ChatResponseTurn2)[]>;
44+
getChatHistory(): (ChatRequestTurn2 | ChatResponseTurn2)[];
3945
}
4046

4147
export class CopilotCLISession extends DisposableStore implements ICopilotCLISession {
@@ -49,25 +55,49 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
4955

5056
public readonly onDidChangeStatus = this._statusChange.event;
5157

58+
private _permissionRequested?: PermissionRequest;
59+
public get permissionRequested(): PermissionRequest | undefined {
60+
return this._permissionRequested;
61+
}
62+
private readonly _onPermissionRequested = this.add(new EventEmitter<PermissionRequest>());
63+
public readonly onPermissionRequested = this._onPermissionRequested.event;
64+
private _permissionHandler?: PermissionHandler;
65+
private readonly _permissionHandlerSet = this.add(new Emitter<void>());
66+
private _stream?: vscode.ChatResponseStream;
5267
constructor(
68+
private readonly _options: CopilotCLISessionOptions,
5369
private readonly _sdkSession: Session,
54-
private readonly _options: SessionOptions,
55-
private readonly _permissionHandler: CopilotCLIPermissionsHandler,
5670
@ILogService private readonly logService: ILogService,
5771
@IWorkspaceService private readonly workspaceService: IWorkspaceService,
58-
@IToolsService private readonly toolsService: IToolsService,
5972
@ICopilotCLISessionOptionsService private readonly cliSessionOptions: ICopilotCLISessionOptionsService,
6073
) {
6174
super();
6275
this.sessionId = _sdkSession.sessionId;
6376
}
6477

78+
attachStream(stream: vscode.ChatResponseStream): IDisposable {
79+
this._stream = stream;
80+
return toDisposable(() => {
81+
if (this._stream === stream) {
82+
this._stream = undefined;
83+
}
84+
});
85+
}
86+
87+
attachPermissionHandler(handler: PermissionHandler): IDisposable {
88+
this._permissionHandler = handler;
89+
this._permissionHandlerSet.fire();
90+
return toDisposable(() => {
91+
if (this._permissionHandler === handler) {
92+
this._permissionHandler = undefined;
93+
}
94+
});
95+
}
96+
6597
public async handleRequest(
6698
prompt: string,
6799
attachments: Attachment[],
68100
modelId: string | undefined,
69-
stream: vscode.ChatResponseStream,
70-
toolInvocationToken: vscode.ChatParticipantToolToken,
71101
token: vscode.CancellationToken
72102
): Promise<void> {
73103
if (this.isDisposed) {
@@ -88,26 +118,25 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
88118
const editToolIds = new Set<string>();
89119
const editTracker = new ExternalEditTracker();
90120
const editFilesAndToolCallIds = new ResourceMap<string[]>();
91-
disposables.add(this._permissionHandler.onDidRequestPermissions(async (permissionRequest) => {
92-
return await this.requestPermission(permissionRequest, stream, editTracker,
121+
disposables.add(this._options.addPermissionHandler(async (permissionRequest) => {
122+
// Need better API from SDK to correlate file edits in permission requests to tool invocations.
123+
return await this.requestPermission(permissionRequest, editTracker,
93124
(file: Uri) => {
94125
const ids = editFilesAndToolCallIds.get(file);
95126
return ids?.shift();
96127
},
97-
toolInvocationToken,
98-
this._options.workingDirectory
128+
this._options.toSessionOptions().workingDirectory,
129+
token
99130
);
100131
}));
101132

102133
try {
103-
const [currentModel,
104-
sessionOptions
105-
] = await Promise.all([
134+
const [currentModel, authInfo] = await Promise.all([
106135
modelId ? this._sdkSession.getSelectedModel() : undefined,
107-
this.cliSessionOptions.createOptions(this._options, this._permissionHandler)
136+
this.cliSessionOptions.createOptions({}).then(opts => opts.toSessionOptions().authInfo)
108137
]);
109-
if (sessionOptions.authInfo) {
110-
this._sdkSession.setAuthInfo(sessionOptions.authInfo);
138+
if (authInfo) {
139+
this._sdkSession.setAuthInfo(authInfo);
111140
}
112141
if (modelId && modelId !== currentModel) {
113142
await this._sdkSession.setSelectedModel(modelId);
@@ -116,7 +145,7 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
116145
disposables.add(toDisposable(this._sdkSession.on('*', (event) => this.logService.trace(`[CopilotCLISession]CopilotCLI Event: ${JSON.stringify(event, null, 2)}`))));
117146
disposables.add(toDisposable(this._sdkSession.on('assistant.message', (event) => {
118147
if (typeof event.data.content === 'string' && event.data.content.length) {
119-
stream.markdown(event.data.content);
148+
this._stream?.markdown(event.data.content);
120149
}
121150
})));
122151
disposables.add(toDisposable(this._sdkSession.on('tool.execution_start', (event) => {
@@ -136,7 +165,7 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
136165
} else {
137166
const responsePart = processToolExecutionStart(event, this._pendingToolInvocations);
138167
if (responsePart instanceof ChatResponseThinkingProgressPart) {
139-
stream.push(responsePart);
168+
this._stream?.push(responsePart);
140169
}
141170
}
142171
this.logService.trace(`[CopilotCLISession] Start Tool ${event.data.toolName || '<unknown>'}`);
@@ -151,7 +180,7 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
151180

152181
const responsePart = processToolExecutionComplete(event, this._pendingToolInvocations);
153182
if (responsePart && !(responsePart instanceof ChatResponseThinkingProgressPart)) {
154-
stream.push(responsePart);
183+
this._stream?.push(responsePart);
155184
}
156185

157186
const toolName = toolNames.get(event.data.toolCallId) || '<unknown>';
@@ -163,7 +192,7 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
163192
})));
164193
disposables.add(toDisposable(this._sdkSession.on('session.error', (event) => {
165194
this.logService.error(`[CopilotCLISession]CopilotCLI error: (${event.data.errorType}), ${event.data.message}`);
166-
stream.markdown(`\n\n❌ Error: (${event.data.errorType}) ${event.data.message}`);
195+
this._stream?.markdown(`\n\n❌ Error: (${event.data.errorType}) ${event.data.message}`);
167196
})));
168197

169198
await this._sdkSession.send({ prompt, attachments, abortController });
@@ -175,7 +204,7 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
175204
this._status = ChatSessionStatus.Failed;
176205
this._statusChange.fire(this._status);
177206
this.logService.error(`[CopilotCLISession] Invoking session (error) ${this.sessionId}`, error);
178-
stream.markdown(`\n\n❌ Error: ${error instanceof Error ? error.message : String(error)}`);
207+
this._stream?.markdown(`\n\n❌ Error: ${error instanceof Error ? error.message : String(error)}`);
179208
} finally {
180209
disposables.dispose();
181210
}
@@ -196,18 +225,17 @@ export class CopilotCLISession extends DisposableStore implements ICopilotCLISes
196225
return this._sdkSession.getSelectedModel();
197226
}
198227

199-
public async getChatHistory(): Promise<(ChatRequestTurn2 | ChatResponseTurn2)[]> {
200-
const events = await this._sdkSession.getEvents();
228+
public getChatHistory(): (ChatRequestTurn2 | ChatResponseTurn2)[] {
229+
const events = this._sdkSession.getEvents();
201230
return buildChatHistoryFromEvents(events);
202231
}
203232

204233
private async requestPermission(
205234
permissionRequest: PermissionRequest,
206-
stream: vscode.ChatResponseStream,
207235
editTracker: ExternalEditTracker,
208236
getEditKeyForFile: (file: Uri) => string | undefined,
209-
toolInvocationToken: vscode.ChatParticipantToolToken,
210-
workingDirectory?: string
237+
workingDirectory: string | undefined,
238+
token: vscode.CancellationToken
211239
): Promise<{ kind: 'approved' } | { kind: 'denied-interactively-by-user' }> {
212240
if (permissionRequest.kind === 'read') {
213241
// 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
239267
}
240268

241269
try {
242-
const { tool, input } = getConfirmationToolParams(permissionRequest);
243-
const result = await this.toolsService.invokeTool(tool,
244-
{ input, toolInvocationToken },
245-
CancellationToken.None);
270+
const permissionHandler = await this.waitForPermissionHandler(permissionRequest);
271+
if (!permissionHandler) {
272+
this.logService.warn(`[CopilotCLISession] No permission handler registered, denying request for ${permissionRequest.kind} permission.`);
273+
return { kind: 'denied-interactively-by-user' };
274+
}
246275

247-
const firstResultPart = result.content.at(0);
248-
if (firstResultPart instanceof LanguageModelTextPart && firstResultPart.value === 'yes') {
276+
if (await permissionHandler(permissionRequest, token)) {
249277
// If we're editing a file, start tracking the edit & wait for core to acknowledge it.
250278
const editFile = permissionRequest.kind === 'write' ? Uri.file(permissionRequest.fileName) : undefined;
251279
const editKey = editFile ? getEditKeyForFile(editFile) : undefined;
252-
if (editFile && editKey) {
280+
if (editFile && editKey && this._stream) {
253281
this.logService.trace(`[CopilotCLISession] Starting to track edit for toolCallId ${editKey} & file ${editFile.fsPath}`);
254-
await editTracker.trackEdit(editKey, [editFile], stream);
282+
await editTracker.trackEdit(editKey, [editFile], this._stream);
255283
}
256284
return { kind: 'approved' };
257285
}
258286
} catch (error) {
259287
this.logService.error(`[CopilotCLISession] Permission request error: ${error}`);
288+
} finally {
289+
this._permissionRequested = undefined;
260290
}
261291

262292
return { kind: 'denied-interactively-by-user' };
263293
}
294+
295+
private async waitForPermissionHandler(permissionRequest: PermissionRequest): Promise<PermissionHandler | undefined> {
296+
if (!this._permissionHandler) {
297+
this._permissionRequested = permissionRequest;
298+
this._onPermissionRequested.fire(permissionRequest);
299+
const disposables = this.add(new DisposableStore());
300+
await Event.toPromise(this._permissionHandlerSet.event, disposables);
301+
disposables.dispose();
302+
this._permissionRequested = undefined;
303+
}
304+
return this._permissionHandler;
305+
}
264306
}

0 commit comments

Comments
 (0)