From 70e927ce69c75aea5b2053574333e6525b717ba3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=93lafur=20P=C3=A1ll=20Geirsson?= Date: Wed, 31 Jan 2024 21:38:29 +0100 Subject: [PATCH] Agent: handle document URIs with exclamation marks (#2983) Previously, the agent threw an error for `textDocument/*` notificattions when the URI included an exclamation mark. The error was thrown to preserve an important variant that we can always parse and render a URI back into the original string ```ts vscode.Uri.parse(uri).toString() === uri ``` This variant was getting violated when the JB plugin sent document URIs that included an exclamation mark in them, for example: ``` file:///Users/com.jetbrains/ideaIC-2022.1-sources.jar!/com/intellij/RequiresBackgroundThread.java ``` The `!` got encoded as `%21`. This PR "fixes" the problem by mutating the `uri: string` property to get the `Uri.toString()` formatted URI. t's not a super satisfying solution, it would be more desirable if we can get `vscode.Uri.parse(..).toString()` to be a safe roundtrip. However, this change gets the most important job done: we preserve the invariant. ## Test plan See updated tests. --- agent/src/AgentWorkspaceDocuments.ts | 4 ++-- agent/src/bfg/BfgRetriever.test.ts | 3 +-- vscode/src/jsonrpc/TextDocumentWithUri.test.ts | 14 ++++++++++++++ vscode/src/jsonrpc/TextDocumentWithUri.ts | 12 ++++++++---- 4 files changed, 25 insertions(+), 8 deletions(-) create mode 100644 vscode/src/jsonrpc/TextDocumentWithUri.test.ts diff --git a/agent/src/AgentWorkspaceDocuments.ts b/agent/src/AgentWorkspaceDocuments.ts index 3de8c276974..0beed3b2ba4 100644 --- a/agent/src/AgentWorkspaceDocuments.ts +++ b/agent/src/AgentWorkspaceDocuments.ts @@ -29,7 +29,7 @@ export class AgentWorkspaceDocuments implements vscode_shim.WorkspaceDocuments { public activeDocumentFilePath: vscode.Uri | null = null public openUri(uri: vscode.Uri): AgentTextDocument { - return this.loadedDocument(new ProtocolTextDocumentWithUri(uri)) + return this.loadedDocument(ProtocolTextDocumentWithUri.from(uri)) } public loadedDocument(document: ProtocolTextDocumentWithUri): AgentTextDocument { const fromCache = this.agentDocuments.get(document.underlying.uri) @@ -122,7 +122,7 @@ export class AgentWorkspaceDocuments implements vscode_shim.WorkspaceDocuments { } public async openTextDocument(uri: vscode.Uri): Promise { - const document = new ProtocolTextDocumentWithUri(uri) + const document = ProtocolTextDocumentWithUri.from(uri) if (!this.agentDocuments.has(document.underlying.uri)) { // Read the file content from disk if the user hasn't opened this file before. const buffer = await fspromises.readFile(uri.fsPath, 'utf8') diff --git a/agent/src/bfg/BfgRetriever.test.ts b/agent/src/bfg/BfgRetriever.test.ts index 7b2380c2876..7606652623f 100644 --- a/agent/src/bfg/BfgRetriever.test.ts +++ b/agent/src/bfg/BfgRetriever.test.ts @@ -12,7 +12,6 @@ import * as vscode from 'vscode' import { BfgRetriever } from '../../../vscode/src/completions/context/retrievers/bfg/bfg-retriever' import { getCurrentDocContext } from '../../../vscode/src/completions/get-current-doc-context' import { initTreeSitterParser } from '../../../vscode/src/completions/test-helpers' -import { ProtocolTextDocumentWithUri } from '../../../vscode/src/jsonrpc/TextDocumentWithUri' import { initializeVscodeExtension, newEmbeddedAgentClient } from '../agent' import * as vscode_shim from '../vscode-shim' @@ -96,7 +95,7 @@ describe('BfgRetriever', async () => { const bfg = new BfgRetriever(extensionContext as vscode.ExtensionContext) - const document = agent.workspace.getDocument(new ProtocolTextDocumentWithUri(uri).uri)! + const document = agent.workspace.getDocument(uri)! assert(document.getText().length > 0) const offset = content.indexOf(CURSOR) assert(offset >= 0, content) diff --git a/vscode/src/jsonrpc/TextDocumentWithUri.test.ts b/vscode/src/jsonrpc/TextDocumentWithUri.test.ts new file mode 100644 index 00000000000..f070b504756 --- /dev/null +++ b/vscode/src/jsonrpc/TextDocumentWithUri.test.ts @@ -0,0 +1,14 @@ +import { describe, expect, it } from 'vitest' +import { ProtocolTextDocumentWithUri } from './TextDocumentWithUri' + +describe('TextDocumentWithUri', () => { + it('handles URIs with exclamation marks', () => { + const uri = + 'file:///Users/com.jetbrains/ideaIC-2022.1-sources.jar!/com/intellij/RequiresBackgroundThread.java' + const textDocument = ProtocolTextDocumentWithUri.fromDocument({ uri }) + expect(textDocument.uri.toString()).toStrictEqual(textDocument.underlying.uri) + expect(textDocument.uri.toString()).toStrictEqual( + 'file:///Users/com.jetbrains/ideaIC-2022.1-sources.jar%21/com/intellij/RequiresBackgroundThread.java' + ) + }) +}) diff --git a/vscode/src/jsonrpc/TextDocumentWithUri.ts b/vscode/src/jsonrpc/TextDocumentWithUri.ts index b422bdcf60c..d8101f02724 100644 --- a/vscode/src/jsonrpc/TextDocumentWithUri.ts +++ b/vscode/src/jsonrpc/TextDocumentWithUri.ts @@ -1,6 +1,7 @@ import * as vscode from 'vscode' import type { Range, ProtocolTextDocument } from './agent-protocol' +import { logDebug } from '../log' /** * Wrapper around `ProtocolTextDocument` that also contains a parsed vscode.Uri. @@ -10,15 +11,18 @@ import type { Range, ProtocolTextDocument } from './agent-protocol' */ export class ProtocolTextDocumentWithUri { public underlying: ProtocolTextDocument - constructor( + private constructor( public readonly uri: vscode.Uri, underlying?: ProtocolTextDocument ) { this.underlying = underlying ?? { uri: uri.toString() } if (this.underlying.uri !== uri.toString()) { - throw new Error( - `ProtocolTextDocumentWithUri invariant violation: ${this.uri} (this.uri) !== ${this.underlying.uri} (this.underlying.uri)` + logDebug( + 'ProtocolTextDocumentWithUri', + 'correcting invariant violation', + `${this.uri} (this.uri) !== ${this.underlying.uri} (this.underlying.uri)` ) + this.underlying.uri = uri.toString() } } @@ -34,7 +38,7 @@ export class ProtocolTextDocumentWithUri { public static from( uri: vscode.Uri, - document: Partial + document?: Partial ): ProtocolTextDocumentWithUri { return new ProtocolTextDocumentWithUri(uri, { ...document, uri: uri.toString() }) }