diff --git a/agent/src/TestClient.ts b/agent/src/TestClient.ts
index 283c8e4105..73cba4d4aa 100644
--- a/agent/src/TestClient.ts
+++ b/agent/src/TestClient.ts
@@ -178,6 +178,7 @@ export class TestClient extends MessageHandler {
     public readonly name: string
     public workspace = new AgentWorkspaceDocuments()
     public workspaceEditParams: WorkspaceEditParams[] = []
+    public textDocumentEditParams: TextDocumentEditParams[] = []
 
     get serverEndpoint(): string {
         return this.params.credentials.serverEndpoint
@@ -330,6 +331,7 @@ export class TestClient extends MessageHandler {
             return Promise.resolve(true)
         })
         this.registerRequest('textDocument/edit', params => {
+            this.textDocumentEditParams.push(params)
             return Promise.resolve(this.editDocument(params).success)
         })
         this.registerRequest('textDocument/show', () => {
diff --git a/agent/src/index.test.ts b/agent/src/index.test.ts
index 8f0e3057d8..15493c88f5 100644
--- a/agent/src/index.test.ts
+++ b/agent/src/index.test.ts
@@ -951,8 +951,9 @@ describe('Agent', () => {
                 .allUris()
                 .filter(uri => vscode.Uri.parse(uri).scheme === 'untitled')
             expect(untitledDocuments).toHaveLength(2)
-            const [untitledDocument] = untitledDocuments
-            const testDocument = client.workspace.getDocument(vscode.Uri.parse(untitledDocument))
+            const untitledDocument = untitledDocuments.find(d => d.endsWith('trickyLogic.test.ts'))
+            expect(untitledDocument).toBeDefined()
+            const testDocument = client.workspace.getDocument(vscode.Uri.parse(untitledDocument ?? ''))
             expect(trimEndOfLine(testDocument?.getText())).toMatchInlineSnapshot(
                 `
               "import { expect } from 'vitest'
@@ -984,9 +985,7 @@ describe('Agent', () => {
                 explainPollyError
             )
 
-            // Just to make sure the edit happened via `workspace/edit` instead
-            // of `textDocument/edit`.
-            expect(client.workspaceEditParams).toHaveLength(1)
+            expect(client.textDocumentEditParams).toHaveLength(1)
         }, 30_000)
 
         describe('Edit code', () => {
diff --git a/agent/src/vscode-shim.ts b/agent/src/vscode-shim.ts
index ff7464ac5f..52d676d885 100644
--- a/agent/src/vscode-shim.ts
+++ b/agent/src/vscode-shim.ts
@@ -4,7 +4,7 @@ import path from 'node:path'
 import * as uuid from 'uuid'
 import type * as vscode from 'vscode'
 
-import { logDebug, logError, setClientNameVersion } from '@sourcegraph/cody-shared'
+import { extensionForLanguage, logDebug, logError, setClientNameVersion } from '@sourcegraph/cody-shared'
 
 // <VERY IMPORTANT - PLEASE READ>
 // This file must not import any module that transitively imports from 'vscode'.
@@ -292,10 +292,12 @@ const _workspace: typeof vscode.workspace = {
             throw new Error('workspaceDocuments is uninitialized')
         }
 
-        const uri = toUri(uriOrString)
-        if (uri) {
-            if (uri.scheme === 'untitled') await openUntitledDocument(uri)
-            return workspaceDocuments.openTextDocument(uri)
+        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)}`)
@@ -452,20 +454,32 @@ 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 | undefined {
+): { uri: Uri; shouldOpenInClient: boolean } | undefined {
     if (typeof uriOrString === 'string') {
-        return Uri.file(uriOrString)
+        return { uri: Uri.file(uriOrString), shouldOpenInClient: true }
     }
     if (uriOrString instanceof Uri) {
-        return uriOrString
+        return { uri: uriOrString, shouldOpenInClient: true }
     }
     if (
         typeof uriOrString === 'object' &&
         ((uriOrString as any)?.language || (uriOrString as any)?.content)
     ) {
-        return Uri.from({ scheme: 'untitled', path: `${uuid.v4()}` })
+        const language = (uriOrString as any)?.language ?? ''
+        const extension = extensionForLanguage(language) ?? language
+        return {
+            uri: Uri.from({
+                scheme: 'untitled',
+                path: `${uuid.v4()}.${extension}`,
+            }),
+            shouldOpenInClient: false,
+        }
     }
     return
 }
diff --git a/vscode/src/edit/provider.ts b/vscode/src/edit/provider.ts
index daef3274cc..83c528a39a 100644
--- a/vscode/src/edit/provider.ts
+++ b/vscode/src/edit/provider.ts
@@ -329,13 +329,10 @@ export class EditProvider {
             if (!fileIsFound) {
                 newFileUri = newFileUri.with({ scheme: 'untitled' })
             }
-            this.insertionPromise = this.config.controller.didReceiveNewFileRequest(task.id, newFileUri)
             try {
-                await this.insertionPromise
+                await this.config.controller.didReceiveNewFileRequest(task.id, newFileUri)
             } catch (error) {
                 this.handleError(new Error('Cody failed to generate unit tests', { cause: error }))
-            } finally {
-                this.insertionPromise = null
             }
         }
     }