feat: support cancellation of RPC requests

lean4-infoview-api/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@leanprover/infoview-api",
-    "version": "0.10.0",
+    "version": "0.11.0",
     "description": "Types and API for @leanprover/infoview.",
     "scripts": {
         "watch": "tsc --watch",

lean4-infoview-api/src/infoviewApi.ts

@@ -45,12 +45,31 @@ export interface InfoviewConfig {
  */
 export type TextInsertKind = 'here' | 'above'
 
+export interface ClientRequestOptions {
+    /**
+     * Can be used to cancel the request.
+     *
+     * This does not immediately reject the promise;
+     * instead, an LSP `$/cancelRequest` notification is sent to the Lean server
+     * The request handler can choose to ignore it,
+     * to return a partial result,
+     * or to produce a `RequestCancelled` error.
+     *
+     * Does nothing if the promise has already been resolved or rejected.
+     */
+    abortSignal?: AbortSignal
+    /**
+     * If `true`, the request will be automatically cancelled when the infoview is closed.
+     */
+    autoCancel?: boolean
+}
+
 /** Interface that the InfoView WebView uses to talk to the hosting editor. */
 export interface EditorApi {
     saveConfig(config: InfoviewConfig): Promise<any>
 
     /** Make a request to the LSP server. */
-    sendClientRequest(uri: string, method: string, params: any): Promise<any>
+    sendClientRequest(uri: string, method: string, params: any, options?: ClientRequestOptions): Promise<any>
     /** Send a notification to the LSP server. */
     sendClientNotification(uri: string, method: string, params: any): Promise<void>
 

lean4-infoview-api/src/lspTypes.ts

@@ -128,6 +128,8 @@ export interface RpcError {
     message: string
 }
 
-export function isRpcError(x: any): x is RpcError {
-    return !!(x?.code && x?.message)
+export function isRpcError(x: unknown): x is RpcError {
+    if (x == null || typeof x !== 'object') return false
+    const obj = x as Record<string, unknown>
+    return typeof obj.code === 'number' && typeof obj.message === 'string'
 }

lean4-infoview-api/src/rpcSessions.ts

@@ -1,4 +1,5 @@
 import type { DocumentUri, Position, TextDocumentPositionParams } from 'vscode-languageserver-protocol'
+import { ClientRequestOptions } from './infoviewApi'
 import { RpcCallParams, RpcErrorCode, RpcPtr, RpcReleaseParams } from './lspTypes'
 
 /**
@@ -23,7 +24,7 @@ export interface RpcServerIface {
     /** Closes an RPC session created with {@link createRpcSession}. */
     closeRpcSession(sessionId: string): void
     /** Sends an RPC call to the Lean server. */
-    call(request: RpcCallParams): Promise<any>
+    call(request: RpcCallParams, options?: ClientRequestOptions): Promise<any>
     /** Sends an RPC reference release notification to the server. */
     release(request: RpcReleaseParams): void
 }
@@ -45,9 +46,10 @@ export interface RpcSessionAtPos {
      *
      * @param method fully qualified name of the method to call.
      * @param params arguments to the invoked method.
+     * @param options additional configuration for the request.
      * @returns a promise that resolves to the returned value, or an error in case the call fails.
      */
-    call<T, S>(method: string, params: T): Promise<S>
+    call<T, S>(method: string, params: T, options?: ClientRequestOptions): Promise<S>
 }
 
 /**
@@ -181,14 +183,15 @@ class RpcSessionForFile {
      * @param position within the file identified by {@link uri}, used to resolve the set of available RPC methods.
      * @param method fully qualified name of the method to call.
      * @param params arguments to the invoked method.
+     * @param options additional configuration for the request.
      * @returns a promise that resolves to the returned value, or to an error in case the call fails.
      */
-    async call(position: Position, method: string, params: any): Promise<any> {
+    async call(position: Position, method: string, params: any, options?: ClientRequestOptions): Promise<any> {
         const sessionId = await this.sessionId
         if (this.failed) throw this.failed
         const tdpp: TextDocumentPositionParams = { position, textDocument: { uri: this.uri } }
         try {
-            const result = await this.sessions.iface.call({ method, params, sessionId, ...tdpp })
+            const result = await this.sessions.iface.call({ method, params, sessionId, ...tdpp }, options)
             this.registerRefs(result)
             // HACK: most of our types are `T | undefined` so try to return something matching that interface
             if (result === null) return undefined
@@ -216,7 +219,9 @@ class RpcSessionForFile {
         // As JS tradition dictates, we use stringification for deep comparison of `Position`s in a `Map`.
         const posStr = `${position.line}:${position.character}`
         if (this.sessionsAtPos.has(posStr)) return this.sessionsAtPos.get(posStr) as RpcSessionAtPos
-        const atPos: RpcSessionAtPos = { call: (method, params) => this.call(position, method, params) }
+        const atPos: RpcSessionAtPos = {
+            call: (method, params, options) => this.call(position, method, params, options),
+        }
         this.sessionsAtPos.set(posStr, atPos)
         return atPos
     }

lean4-infoview/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@leanprover/infoview",
-    "version": "0.10.0",
+    "version": "0.11.0",
     "description": "An interactive display for the Lean 4 theorem prover.",
     "scripts": {
         "watch": "rollup --config --environment NODE_ENV:development --watch",
@@ -49,7 +49,7 @@
         "typescript": "^5.4.5"
     },
     "dependencies": {
-        "@leanprover/infoview-api": "~0.10.0",
+        "@leanprover/infoview-api": "~0.11.0",
         "@vscode/codicons": "^0.0.40",
         "@vscode-elements/react-elements": "^0.5.0",
         "es-module-lexer": "^1.5.4",

lean4-infoview/src/infoview/rpcSessions.tsx

@@ -21,8 +21,8 @@ export function WithRpcSessions({ children }: { children: React.ReactNode }) {
             new RpcSessions({
                 createRpcSession: (uri: DocumentUri) => ec.api.createRpcSession(uri),
                 closeRpcSession: (uri: DocumentUri) => ec.api.closeRpcSession(uri),
-                call: (params: RpcCallParams) =>
-                    ec.api.sendClientRequest(params.textDocument.uri, '$/lean/rpc/call', params),
+                call: (params: RpcCallParams, options) =>
+                    ec.api.sendClientRequest(params.textDocument.uri, '$/lean/rpc/call', params, options),
                 release: (params: RpcReleaseParams) =>
                     void ec.api.sendClientNotification(params.uri, '$/lean/rpc/release', params),
             }),

vscode-lean4/package.json

@@ -1761,8 +1761,8 @@
         "test": "node ./out/test/suite/runTest.js"
     },
     "dependencies": {
-        "@leanprover/infoview": "~0.10.0",
-        "@leanprover/infoview-api": "~0.10.0",
+        "@leanprover/infoview": "~0.11.0",
+        "@leanprover/infoview-api": "~0.11.0",
         "@leanprover/unicode-input": "~0.1.8",
         "@leanprover/unicode-input-component": "~0.1.8",
         "@vscode/codicons": "^0.0.36",

vscode-lean4/src/infoview.ts

@@ -1,5 +1,4 @@
 import {
-    EditorApi,
     InfoviewApi,
     InfoviewConfig,
     LeanFileProgressParams,
@@ -12,6 +11,7 @@ import {
 } from '@leanprover/infoview-api'
 import { join } from 'path'
 import {
+    CancellationTokenSource,
     commands,
     ConfigurationTarget,
     Diagnostic,
@@ -50,7 +50,7 @@ import {
     prodOrDev,
 } from './config'
 import { LeanClient } from './leanclient'
-import { Rpc } from './rpc'
+import { ClientRequestRpcOptions, EditorRpcApi, Rpc } from './rpc'
 import { LeanClientProvider } from './utils/clientProvider'
 import { c2pConverter, LeanPublishDiagnosticsParams, p2cConverter } from './utils/converters'
 import { ExtUri, parseExtUri, toExtUri } from './utils/exturi'
@@ -97,8 +97,8 @@ class RpcSessionAtPos implements Disposable {
 export class InfoProvider implements Disposable {
     /** Instance of the panel, if it is open. Otherwise `undefined`. */
     private webviewPanel?: WebviewPanel & { rpc: Rpc; api: InfoviewApi }
+    /** The InfoProvider's subscriptions, to be cleaned up when it is disposed. */
     private subscriptions: Disposable[] = []
-    private clientSubscriptions: Disposable[] = []
 
     private stylesheet: string = ''
     private autoOpened: boolean = false
@@ -115,6 +115,16 @@ export class InfoProvider implements Disposable {
     // the key is the uri of the file who's worker has failed.
     private workersFailed: Map<string, ServerStoppedReason> = new Map()
 
+    /**
+     * The ID to assign to the next client request made by the infoview
+     * (see {@link editorApi.startClientRequest}).
+     * Only used for cancellation. Unrelated to LSP request IDs. */
+    private freshClientRequestId: number = 0
+    /**
+     * Maps in-flight client request IDs (as in {@link freshClientRequestId})
+     * to tuples [Response, Infoview subscription (should be disposed when infoview closes), Cancellation token]. */
+    private clientRequests: Map<number, [Promise<any>, Disposable, CancellationTokenSource]> = new Map()
+
     private subscribeDidChangeNotification(client: LeanClient, method: string) {
         const h = client.didChange(params => {
             void this.webviewPanel?.api.sentClientNotification(method, params)
@@ -144,7 +154,7 @@ export class InfoProvider implements Disposable {
         return h
     }
 
-    private editorApi: EditorApi = {
+    private editorApi: EditorRpcApi = {
         saveConfig: async (config: InfoviewConfig) => {
             await workspace
                 .getConfiguration('lean4.infoview')
@@ -186,18 +196,28 @@ export class InfoProvider implements Disposable {
                 .getConfiguration('lean4.infoview')
                 .update('messageOrder', config.messageOrder, ConfigurationTarget.Global)
         },
-        sendClientRequest: async (uri: string, method: string, params: any): Promise<any> => {
+        startClientRequest: async (
+            uri: string,
+            method: string,
+            params: any,
+            options?: ClientRequestRpcOptions,
+        ): Promise<number> => {
             const extUri = parseExtUri(uri)
             if (extUri === undefined) {
                 throw Error(`Unexpected URI scheme: ${Uri.parse(uri).scheme}`)
             }
 
             const client = this.clientProvider.findClient(extUri)
             if (client) {
-                try {
-                    const result = await client.sendRequest(method, params)
-                    return result
-                } catch (ex) {
+                const tk = new CancellationTokenSource()
+                const sub: Disposable = {
+                    dispose() {
+                        if (options?.autoCancel) tk.cancel()
+                    },
+                }
+                const id = this.freshClientRequestId
+                this.freshClientRequestId += 1
+                const promise = client.sendRequest(method, params, tk.token).catch(async ex => {
                     if (ex.code === RpcErrorCode.WorkerCrashed) {
                         // ex codes related with worker exited or crashed
                         logger.log(`[InfoProvider]The Lean Server has stopped processing this file: ${ex.message}`)
@@ -207,10 +227,34 @@ export class InfoProvider implements Disposable {
                         })
                     }
                     throw ex
-                }
+                })
+                this.clientRequests.set(id, [promise, sub, tk])
+                return id
             }
             throw Error('No active Lean client.')
         },
+        awaitClientRequest: async (id: number): Promise<any> => {
+            const p = this.clientRequests.get(id)
+            if (p) {
+                // NOTE: `await` finishes first (or throws),
+                // then the entry is deleted,
+                // and then the function returns.
+                try {
+                    return await p[0]
+                } finally {
+                    this.clientRequests.delete(id)
+                }
+            }
+            // NOTE: we rely on details of `editorApiOfRpc` for correctness:
+            // we assume that `awaitClientRequest` is called exactly once
+            // regardless of whether cancellation occurs,
+            // so this error should never be thrown.
+            throw Error(`Internal error: invalid client request ID '${id}'`)
+        },
+        cancelClientRequest: async (id: number): Promise<void> => {
+            const p = this.clientRequests.get(id)
+            if (p) p[2].cancel()
+        },
         sendClientNotification: async (uri: string, method: string, params: any): Promise<void> => {
             const extUri = parseExtUri(uri)
             if (extUri === undefined) {
@@ -370,19 +414,16 @@ export class InfoProvider implements Disposable {
     ) {
         this.updateStylesheet()
 
-        clientProvider.clientAdded(client => {
-            void this.onClientAdded(client)
-        })
-
-        clientProvider.clientRemoved(client => {
-            void this.onClientRemoved(client)
-        })
-
-        clientProvider.clientStopped(([client, activeClient, reason]) => {
-            void this.onActiveClientStopped(client, activeClient, reason)
-        })
-
         this.subscriptions.push(
+            clientProvider.clientAdded(client => {
+                void this.onClientAdded(client)
+            }),
+            clientProvider.clientRemoved(client => {
+                void this.onClientRemoved(client)
+            }),
+            clientProvider.clientStopped(([client, activeClient, reason]) => {
+                void this.onActiveClientStopped(client, activeClient, reason)
+            }),
             lean.onDidChangeActiveLeanEditor(() => this.sendPosition()),
             lean.onDidChangeLeanEditorSelection(() => this.sendPosition()),
             workspace.onDidChangeConfiguration(async _e => {
@@ -605,7 +646,7 @@ export class InfoProvider implements Disposable {
     private async onClientAdded(client: LeanClient) {
         logger.log(`[InfoProvider] Adding client for workspace: ${client.getClientFolder()}`)
 
-        this.clientSubscriptions.push(
+        this.subscriptions.push(
             client.restarted(async () => {
                 logger.log('[InfoProvider] got client restarted event')
                 // This event is triggered both the first time the server starts
@@ -654,7 +695,7 @@ export class InfoProvider implements Disposable {
     }
 
     onClientRemoved(client: LeanClient) {
-        // todo: remove subscriptions for this client...
+        // NOTE: we could index subscriptions made in `onClientAdded` by the client, and remove here
     }
 
     async onActiveClientStopped(client: LeanClient, activeClient: boolean, reason: ServerStoppedReason) {
@@ -677,17 +718,11 @@ export class InfoProvider implements Disposable {
         client.showRestartMessage()
     }
 
+    // Called when the extension is deactivated.
     dispose(): void {
-        // active client is changing.
-        this.clearNotificationHandlers()
-        this.clearRpcSessions(null)
+        // Calls `webviewPanel.onDidDispose` to cleanup the infoview's subscriptions.
         this.webviewPanel?.dispose()
-        for (const s of this.clientSubscriptions) {
-            s.dispose()
-        }
-        for (const s of this.subscriptions) {
-            s.dispose()
-        }
+        for (const s of this.subscriptions) s.dispose()
     }
 
     isOpen(): boolean {
@@ -819,6 +854,8 @@ export class InfoProvider implements Disposable {
                 this.webviewPanel = undefined
                 this.clearNotificationHandlers()
                 this.clearRpcSessions(null) // should be after `webviewPanel = undefined`
+                for (const [_1, d, _2] of this.clientRequests.values()) d.dispose()
+                this.clientRequests = new Map()
             })
             this.webviewPanel = webviewPanel
             webviewPanel.webview.html = this.initialHtml()