Allow clients to redefine untitled files protocol during their creation

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentClient.kt

@@ -15,7 +15,7 @@ interface CodyAgentClient {
   @JsonRequest("textDocument/edit")
   fun textDocument_edit(params: TextDocumentEditParams): CompletableFuture<Boolean>
   @JsonRequest("textDocument/openUntitledDocument")
-  fun textDocument_openUntitledDocument(params: UntitledTextDocument): CompletableFuture<Boolean>
+  fun textDocument_openUntitledDocument(params: UntitledTextDocument): CompletableFuture<ProtocolTextDocument?>
   @JsonRequest("textDocument/show")
   fun textDocument_show(params: TextDocument_ShowParams): CompletableFuture<Boolean>
   @JsonRequest("workspace/edit")

agent/src/TestClient.ts

@@ -340,9 +340,9 @@ export class TestClient extends MessageHandler {
             return result
         })
         this.registerRequest('textDocument/openUntitledDocument', params => {
-            this.workspace.loadDocument(ProtocolTextDocumentWithUri.fromDocument(params))
+            const doc = this.workspace.loadDocument(ProtocolTextDocumentWithUri.fromDocument(params))
             this.notify('textDocument/didOpen', params)
-            return Promise.resolve(true)
+            return Promise.resolve(doc.protocolDocument.underlying)
         })
         this.registerRequest('textDocument/edit', async params => {
             this.textDocumentEditParams.push(params)

agent/src/agent.ts

@@ -487,7 +487,7 @@ export class Agent extends MessageHandler implements ExtensionClient {
 
         this.registerNotification('textDocument/didSave', async params => {
             const uri = vscode.Uri.parse(params.uri)
-            const document = await vscode.workspace.openTextDocument(uri)
+            const document = await this.workspace.openTextDocument(uri)
             vscode_shim.onDidSaveTextDocument.fire(document)
         })
 
@@ -1319,6 +1319,27 @@ export class Agent extends MessageHandler implements ExtensionClient {
         return this.fixups
     }
 
+    public openNewDocument = async (
+        _: typeof vscode.workspace,
+        uri: vscode.Uri
+    ): Promise<vscode.TextDocument | undefined> => {
+        if (uri.scheme !== 'untitled') {
+            return vscode_shim.workspace.openTextDocument(uri)
+        }
+
+        if (this.clientInfo?.capabilities?.untitledDocuments !== 'enabled') {
+            const errorMessage =
+                'Client does not support untitled documents. To fix this problem, set `untitledDocuments: "enabled"` in client capabilities'
+            logError('Agent', 'unsupported operation', errorMessage)
+            throw new Error(errorMessage)
+        }
+
+        const result = await this.request('textDocument/openUntitledDocument', {
+            uri: uri.toString(),
+        })
+        return result ? vscode_shim.workspace.openTextDocument(result.uri) : undefined
+    }
+
     private maybeExtension: ExtensionObjects | undefined
 
     public async provide(extension: ExtensionObjects): Promise<vscode.Disposable> {

agent/src/vscode-shim.ts

@@ -301,16 +301,14 @@ const _workspace: typeof vscode.workspace = {
             throw new Error('workspaceDocuments is uninitialized')
         }
 
-        const result = toUri(uriOrString)
-        if (result) {
-            if (result.uri.scheme === 'untitled' && result.shouldOpenInClient) {
-                await openUntitledDocument(result.uri)
-            }
-            return workspaceDocuments.openTextDocument(result.uri)
-        }
-        return Promise.reject(
-            new Error(`workspace.openTextDocument:unsupported argument ${JSON.stringify(uriOrString)}`)
-        )
+        const uri = toUri(uriOrString)
+        return uri
+            ? workspaceDocuments.openTextDocument(uri)
+            : Promise.reject(
+                  new Error(
+                      `workspace.openTextDocument: unsupported argument ${JSON.stringify(uriOrString)}`
+                  )
+              )
     },
     workspaceFolders,
     getWorkspaceFolder: () => {
@@ -463,56 +461,27 @@ const defaultTreeView: vscode.TreeView<any> = {
     title: undefined,
 }
 
-/**
- * @returns An object with a URI and a boolean indicating whether the URI should be opened in the client.
- * This object with UUID path is used only when we want to create in-memory temp files, and those we do not want to send to the clients.
- */
 function toUri(
     uriOrString: string | vscode.Uri | { language?: string; content?: string } | undefined
-): { uri: Uri; shouldOpenInClient: boolean } | undefined {
+): Uri | undefined {
     if (typeof uriOrString === 'string') {
-        return { uri: Uri.file(uriOrString), shouldOpenInClient: true }
+        return Uri.parse(uriOrString)
     }
     if (uriOrString instanceof Uri) {
-        return { uri: uriOrString, shouldOpenInClient: true }
+        return uriOrString
     }
     if (
         typeof uriOrString === 'object' &&
         ((uriOrString as any)?.language || (uriOrString as any)?.content)
     ) {
         const language = (uriOrString as any)?.language ?? ''
         const extension = extensionForLanguage(language) ?? language
-        return {
-            uri: Uri.from({
-                scheme: 'untitled',
-                path: `${uuid.v4()}.${extension}`,
-            }),
-            shouldOpenInClient: false,
-        }
-    }
-    return
-}
-
-async function openUntitledDocument(uri: Uri, content?: string, language?: string) {
-    if (clientInfo?.capabilities?.untitledDocuments !== 'enabled') {
-        const errorMessage =
-            'Client does not support untitled documents. To fix this problem, set `untitledDocuments: "enabled"` in client capabilities'
-        logError('vscode.workspace.openTextDocument', 'unsupported operation', errorMessage)
-        throw new Error(errorMessage)
-    }
-    if (agent) {
-        const result = await agent.request('textDocument/openUntitledDocument', {
-            uri: uri.toString(),
-            content,
-            language,
+        return Uri.from({
+            scheme: 'untitled',
+            path: `${uuid.v4()}.${extension}`,
         })
-
-        if (!result) {
-            throw new Error(
-                `client returned false from textDocument/openUntitledDocument: ${uri.toString()}`
-            )
-        }
     }
+    return
 }
 
 function outputChannel(name: string): vscode.LogOutputChannel {
@@ -705,6 +674,7 @@ const _window: typeof vscode.window = {
                       selection.end.character
                   )
                 : undefined
+
             const result = await agent.request('textDocument/show', {
                 uri,
                 options: {
@@ -715,6 +685,7 @@ const _window: typeof vscode.window = {
             if (!result) {
                 throw new Error(`showTextDocument: client returned false when trying to show URI ${uri}`)
             }
+
             if (!workspaceDocuments) {
                 throw new Error('workspaceDocuments is undefined')
             }

vscode/src/extension-client.ts

@@ -1,4 +1,5 @@
-import type { Disposable } from 'vscode'
+import type { Disposable, TextDocument, Uri } from 'vscode'
+import type vscode from 'vscode'
 import type { EnterpriseContextFactory } from './context/enterprise-context-factory'
 import type { ClientCapabilities } from './jsonrpc/agent-protocol'
 import { FixupCodeLenses } from './non-stop/codelenses/provider'
@@ -31,6 +32,12 @@ export interface ExtensionClient {
      */
     createFixupControlApplicator(fixups: FixupActor & FixupFileCollection): FixupControlApplicator
 
+    /**
+     * Opens a new document, creating appropriate file is required by a protocol.
+     * This method allows client to change the URI, so the caller should inspect returned TextDocument.
+     */
+    openNewDocument(workspace: typeof vscode.workspace, uri: Uri): Thenable<TextDocument | undefined>
+
     get clientName(): string
     get clientVersion(): string
     get capabilities(): ClientCapabilities | undefined
@@ -45,6 +52,7 @@ export function defaultVSCodeExtensionClient(): ExtensionClient {
             dispose: () => {},
         }),
         createFixupControlApplicator: files => new FixupCodeLenses(files),
+        openNewDocument: (workspace, uri) => workspace.openTextDocument(uri),
         clientName: 'vscode',
         clientVersion: version,
         capabilities: undefined,

vscode/src/jsonrpc/agent-protocol.ts

@@ -311,7 +311,7 @@ export type ServerRequests = {
     'window/showMessage': [ShowWindowMessageParams, string | null]
 
     'textDocument/edit': [TextDocumentEditParams, boolean]
-    'textDocument/openUntitledDocument': [UntitledTextDocument, boolean]
+    'textDocument/openUntitledDocument': [UntitledTextDocument, ProtocolTextDocument | undefined | null]
     'textDocument/show': [
         {
             uri: string

vscode/src/non-stop/FixupController.ts

@@ -71,7 +71,7 @@ export class FixupController
 
     constructor(
         private readonly authProvider: AuthProvider,
-        client: ExtensionClient
+        private readonly client: ExtensionClient
     ) {
         this.controlApplicator = client.createFixupControlApplicator(this)
         // Observe file renaming and deletion
@@ -985,16 +985,20 @@ export class FixupController
         }
 
         // append response to new file
-        const doc = await vscode.workspace.openTextDocument(newFileUri)
+        const doc = await this.client.openNewDocument(vscode.workspace, newFileUri)
+        if (!doc) {
+            throw new Error(`Cannot create file for the fixup: ${newFileUri.toString()}`)
+        }
+
         const pos = new vscode.Position(Math.max(doc.lineCount - 1, 0), 0)
         const range = new vscode.Range(pos, pos)
         task.selectionRange = range
         task.insertionPoint = range.start
-        task.fixupFile = this.files.replaceFile(task.fixupFile.uri, newFileUri)
+        task.fixupFile = this.files.replaceFile(task.fixupFile.uri, doc.uri)
 
         // Set original text to empty as we are not replacing original text but appending to file
         task.original = ''
-        task.destinationFile = newFileUri
+        task.destinationFile = doc.uri
 
         // Show the new document before streaming start
         await vscode.window.showTextDocument(doc, {