Skip to content

Commit

Permalink
Agent: add cancelation support
Browse files Browse the repository at this point in the history 
Previously, the agent didn't support canceling autocomplete requests.
This PR adds support for cancellation through the new `$/cancelRequest`
notification, which mirrors the structure of the same notification in
LSP.
  • Loading branch information
@olafurpg
olafurpg committed Aug 22, 2023
1 parent c41cbc0 commit 3efc372
Show file tree
Hide file tree
Showing 5 changed files with 93 additions and 37 deletions.
3 changes: 1 addition & 2 deletions agent/src/agent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -162,13 +162,12 @@ export class Agent extends MessageHandler {
})
return null
})
this.registerRequest('autocomplete/execute', async params => {
this.registerRequest('autocomplete/execute', async (params, token) => {
const provider = await vscode_shim.completionProvider
if (!provider) {
console.log('Completion provider is not initialized')
return { items: [] }
}
const token = new vscode.CancellationTokenSource().token
const document = this.workspace.getDocument(params.filePath)
if (!document) {
console.log('No document found for file path', params.filePath, [...this.workspace.allFilePaths()])
Expand Down
80 changes: 51 additions & 29 deletions agent/src/jsonrpc.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ import { appendFileSync, existsSync, mkdirSync, rmSync } from 'fs'
import { dirname } from 'path'
import { Readable, Writable } from 'stream'

import * as vscode from 'vscode'

import { Notifications, Requests } from './protocol'

// This file is a standalone implementation of JSON-RPC for Node.js
Expand Down Expand Up @@ -201,7 +203,10 @@ class MessageEncoder extends Readable {
}
}

type RequestCallback<M extends RequestMethodName> = (params: ParamsOf<M>) => Promise<ResultOf<M>>
type RequestCallback<M extends RequestMethodName> = (
params: ParamsOf<M>,
cancelToken: vscode.CancellationToken
) => Promise<ResultOf<M>>
type NotificationCallback<M extends NotificationMethodName> = (params: ParamsOf<M>) => void

/**
Expand All @@ -211,6 +216,7 @@ type NotificationCallback<M extends NotificationMethodName> = (params: ParamsOf<
export class MessageHandler {
private id = 0
private requestHandlers: Map<RequestMethodName, RequestCallback<any>> = new Map()
private cancelTokens: Map<Id, vscode.CancellationTokenSource> = new Map()
private notificationHandlers: Map<NotificationMethodName, NotificationCallback<any>> = new Map()
private responseHandlers: Map<Id, (params: any) => void> = new Map()

Expand All @@ -231,31 +237,38 @@ export class MessageHandler {
// Requests have ids and methods
const handler = this.requestHandlers.get(msg.method)
if (handler) {
handler(msg.params).then(
result => {
const data: ResponseMessage<any> = {
jsonrpc: '2.0',
id: msg.id,
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
result,
}
this.messageEncoder.send(data)
},
error => {
const message = error instanceof Error ? error.message : `${error}`
const stack = error instanceof Error ? `\n${error.stack}` : ''
const data: ResponseMessage<any> = {
jsonrpc: '2.0',
id: msg.id,
error: {
code: ErrorCode.InternalError,
message,
data: JSON.stringify({ error, stack }),
},
const cancelToken = new vscode.CancellationTokenSource()
this.cancelTokens.set(msg.id, cancelToken)
handler(msg.params, cancelToken.token)
.then(
result => {
const data: ResponseMessage<any> = {
jsonrpc: '2.0',
id: msg.id,
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
result,
}
this.messageEncoder.send(data)
},
error => {
const message = error instanceof Error ? error.message : `${error}`
const stack = error instanceof Error ? `\n${error.stack}` : ''
const data: ResponseMessage<any> = {
jsonrpc: '2.0',
id: msg.id,
error: {
code: ErrorCode.InternalError,
message,
data: JSON.stringify({ error, stack }),
},
}
this.messageEncoder.send(data)
}
this.messageEncoder.send(data)
}
)
)
.finally(() => {
this.cancelTokens.get(msg.params.id)?.dispose()
this.cancelTokens.delete(msg.params.id)
})
} else {
console.error(`No handler for request with method ${msg.method}`)
}
Expand All @@ -270,11 +283,20 @@ export class MessageHandler {
}
} else if (msg.method) {
// Notifications have methods
const notificationHandler = this.notificationHandlers.get(msg.method)
if (notificationHandler) {
notificationHandler(msg.params)
if (
msg.method === '$/cancelRequest' &&
msg.params &&
(typeof msg.params.id === 'string' || typeof msg.params.id === 'number')
) {
this.cancelTokens.get(msg.params.id)?.cancel()
this.cancelTokens.delete(msg.params.id)
} else {
console.error(`No handler for notification with method ${msg.method}`)
const notificationHandler = this.notificationHandlers.get(msg.method)
if (notificationHandler) {
notificationHandler(msg.params)
} else {
console.error(`No handler for notification with method ${msg.method}`)
}
}
}
})
Expand Down
6 changes: 6 additions & 0 deletions agent/src/protocol.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,8 @@ export type Notifications = {
// Only the 'uri' property is required, other properties are ignored.
'textDocument/didClose': [TextDocument]

'$/cancelRequest': [CancelParams]

// ================
// Server -> Client
// ================
Expand All @@ -96,6 +98,10 @@ export type Notifications = {
'debug/message': [DebugMessage]
}

export interface CancelParams {
id: string | number
}

export interface AutocompleteParams {
filePath: string
position: Position
Expand Down
1 change: 1 addition & 0 deletions agent/src/vscode-shim.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ export {
emptyEvent,
emptyDisposable,
Range,
Location,
Selection,
Position,
Disposable,
Expand Down
40 changes: 34 additions & 6 deletions vscode/src/testutils/mocks.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import type {
Disposable as VSCodeDisposable,
InlineCompletionTriggerKind as VSCodeInlineCompletionTriggerKind,
Location as VSCodeLocation,
Position as VSCodePosition,
Range as VSCodeRange,
} from 'vscode'
Expand Down Expand Up @@ -336,6 +337,21 @@ export class Position implements VSCodePosition {
}
}

export class Location implements VSCodeLocation {
public range: VSCodeRange

constructor(
public readonly uri: vscode_types.Uri,
rangeOrPosition: VSCodeRange | VSCodePosition
) {
if ('line' in rangeOrPosition && 'character' in rangeOrPosition) {
this.range = new Range(rangeOrPosition, rangeOrPosition)
} else {
this.range = rangeOrPosition
}
}
}

export class Range implements VSCodeRange {
public start: Position
public end: Position
Expand Down Expand Up @@ -478,7 +494,7 @@ export class EventEmitter<T> implements vscode_types.EventEmitter<T> {
}
}
dispose(): void {
// throw new Error('Method not implemented.')
this.listeners.clear()
}
}

Expand All @@ -494,14 +510,26 @@ export enum FileType {
SymbolicLink = 64,
}

export class CancellationTokenSource {
public token: unknown

export class CancellationToken implements vscode_types.CancellationToken {
public isCancellationRequested = false
public emitter = new EventEmitter<void>()
constructor() {
this.token = {
onCancellationRequested() {},
this.emitter.event(() => {
this.isCancellationRequested = true
})
}
onCancellationRequested = this.emitter.event
}
export class CancellationTokenSource implements vscode_types.CancellationTokenSource {
public token = new CancellationToken()
cancel(): void {
if (!this.token.isCancellationRequested) {
this.token.emitter.fire()
}
}
dispose(): void {
this.token.emitter.dispose()
}
}

export const vsCodeMocks = {
Expand Down

0 comments on commit 3efc372

Please sign in to comment.