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.
<!-- Required. See
https://sourcegraph.com/docs/dev/background-information/testing_principles.
-->

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<vscode.TextDocument> {
-        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')

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)

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'
+        )
+    })
+})

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<ProtocolTextDocument>
+        document?: Partial<ProtocolTextDocument>
     ): ProtocolTextDocumentWithUri {
         return new ProtocolTextDocumentWithUri(uri, { ...document, uri: uri.toString() })
     }