Remove virtualDocUri() as a footgun

apps/vscode/src/host/hooks.ts

@@ -20,7 +20,7 @@ import { ExtensionHost, HostWebviewPanel, HostStatementRangeProvider, HostHelpTo
 import { CellExecutor, cellExecutorForLanguage, executableLanguages, isKnitrDocument, pythonWithReticulate } from './executors';
 import { ExecuteQueue } from './execute-queue';
 import { MarkdownEngine } from '../markdown/engine';
-import { virtualDoc, virtualDocUri, adjustedPosition, unadjustedRange, withVirtualDocUri } from "../vdoc/vdoc";
+import { virtualDoc, adjustedPosition, unadjustedRange, withVirtualDocUri } from "../vdoc/vdoc";
 import { EmbeddedLanguage } from '../vdoc/languages';
 
 declare global {

apps/vscode/src/lsp/client.ts

@@ -23,6 +23,7 @@ import {
   LocationLink,
   Definition,
   LogOutputChannel,
+  Uri
 } from "vscode";
 import {
   LanguageClient,
@@ -51,7 +52,7 @@ import {
   adjustedPosition,
   unadjustedRange,
   virtualDoc,
-  virtualDocUri,
+  withVirtualDocUri,
 } from "../vdoc/vdoc";
 import { activateVirtualDocEmbeddedContent } from "../vdoc/vdoc-content";
 import { vdocCompletions } from "../vdoc/vdoc-completion";
@@ -225,19 +226,13 @@ function embeddedHoverProvider(engine: MarkdownEngine) {
 
     const vdoc = await virtualDoc(document, position, engine);
     if (vdoc) {
-      // get uri for hover
-      const vdocUri = await virtualDocUri(vdoc, document.uri, "hover");
-
-      // execute hover
-      try {
-        return getHover(vdocUri.uri, vdoc.language, position);
-      } catch (error) {
-        console.log(error);
-      } finally {
-        if (vdocUri.cleanup) {
-          await vdocUri.cleanup();
+      return await withVirtualDocUri(vdoc, document.uri, "hover", async (uri: Uri) => {
+        try {
+          return await getHover(uri, vdoc.language, position);
+        } catch (error) {
+          console.log(error);
         }
-      }
+      });
     }
 
     // default to server delegation
@@ -255,16 +250,13 @@ function embeddedSignatureHelpProvider(engine: MarkdownEngine) {
   ) => {
     const vdoc = await virtualDoc(document, position, engine);
     if (vdoc) {
-      const vdocUri = await virtualDocUri(vdoc, document.uri, "signature");
-      try {
-        return getSignatureHelpHover(vdocUri.uri, vdoc.language, position, context.triggerCharacter);
-      } catch (error) {
-        return undefined;
-      } finally {
-        if (vdocUri.cleanup) {
-          await vdocUri.cleanup();
+      return await withVirtualDocUri(vdoc, document.uri, "signature", async (uri: Uri) => {
+        try {
+          return await getSignatureHelpHover(uri, vdoc.language, position, context.triggerCharacter);
+        } catch (error) {
+          return undefined;
         }
-      }
+      });
     } else {
       return await next(document, position, context, token);
     }
@@ -280,64 +272,61 @@ function embeddedGoToDefinitionProvider(engine: MarkdownEngine) {
   ): Promise<Definition | LocationLink[] | null | undefined> => {
     const vdoc = await virtualDoc(document, position, engine);
     if (vdoc) {
-      const vdocUri = await virtualDocUri(vdoc, document.uri, "definition");
-      try {
-        const definitions = await commands.executeCommand<
-          ProviderResult<Definition | LocationLink[]>
-        >(
-          "vscode.executeDefinitionProvider",
-          vdocUri.uri,
-          adjustedPosition(vdoc.language, position)
-        );
-        const resolveLocation = (location: Location) => {
-          if (location.uri.toString() === vdocUri.uri.toString()) {
-            return new Location(
-              document.uri,
-              unadjustedRange(vdoc.language, location.range)
-            );
-          } else {
-            return location;
-          }
-        };
-        const resolveLocationLink = (location: LocationLink) => {
-          if (location.targetUri.toString() === vdocUri.uri.toString()) {
-            const locationLink: LocationLink = {
-              targetRange: unadjustedRange(vdoc.language, location.targetRange),
-              originSelectionRange: location.originSelectionRange
-                ? unadjustedRange(vdoc.language, location.originSelectionRange)
-                : undefined,
-              targetSelectionRange: location.targetSelectionRange
-                ? unadjustedRange(vdoc.language, location.targetSelectionRange)
-                : undefined,
-              targetUri: document.uri,
-            };
-            return locationLink;
-          } else {
-            return location;
-          }
-        };
-        if (definitions instanceof Location) {
-          return resolveLocation(definitions);
-        } else if (Array.isArray(definitions) && definitions.length > 0) {
-          if (definitions[0] instanceof Location) {
-            return definitions.map((definition) =>
-              resolveLocation(definition as Location)
-            );
+      return await withVirtualDocUri(vdoc, document.uri, "definition", async (uri: Uri) => {
+        try {
+          const definitions = await commands.executeCommand<
+            ProviderResult<Definition | LocationLink[]>
+          >(
+            "vscode.executeDefinitionProvider",
+            uri,
+            adjustedPosition(vdoc.language, position)
+          );
+          const resolveLocation = (location: Location) => {
+            if (location.uri.toString() === uri.toString()) {
+              return new Location(
+                document.uri,
+                unadjustedRange(vdoc.language, location.range)
+              );
+            } else {
+              return location;
+            }
+          };
+          const resolveLocationLink = (location: LocationLink) => {
+            if (location.targetUri.toString() === uri.toString()) {
+              const locationLink: LocationLink = {
+                targetRange: unadjustedRange(vdoc.language, location.targetRange),
+                originSelectionRange: location.originSelectionRange
+                  ? unadjustedRange(vdoc.language, location.originSelectionRange)
+                  : undefined,
+                targetSelectionRange: location.targetSelectionRange
+                  ? unadjustedRange(vdoc.language, location.targetSelectionRange)
+                  : undefined,
+                targetUri: document.uri,
+              };
+              return locationLink;
+            } else {
+              return location;
+            }
+          };
+          if (definitions instanceof Location) {
+            return resolveLocation(definitions);
+          } else if (Array.isArray(definitions) && definitions.length > 0) {
+            if (definitions[0] instanceof Location) {
+              return definitions.map((definition) =>
+                resolveLocation(definition as Location)
+              );
+            } else {
+              return definitions.map((definition) =>
+                resolveLocationLink(definition as LocationLink)
+              );
+            }
           } else {
-            return definitions.map((definition) =>
-              resolveLocationLink(definition as LocationLink)
-            );
+            return definitions;
           }
-        } else {
-          return definitions;
-        }
-      } catch (error) {
-        return undefined;
-      } finally {
-        if (vdocUri.cleanup) {
-          await vdocUri.cleanup();
+        } catch (error) {
+          return undefined;
         }
-      }
+      });
     } else {
       return await next(document, position, token);
     }

apps/vscode/src/providers/assist/render-assist.ts

@@ -35,7 +35,7 @@ import {
 import { JsonRpcRequestTransport, escapeRegExpCharacters } from "core";
 import { CodeViewCellContext, kCodeViewAssist } from "editor-types";
 import { embeddedLanguage } from "../../vdoc/languages";
-import { virtualDocForCode, virtualDocUri, withVirtualDocUri } from "../../vdoc/vdoc";
+import { virtualDocForCode, withVirtualDocUri } from "../../vdoc/vdoc";
 import { getHover, getSignatureHelpHover } from "../../core/hover";
 import { Hover as LspHover, MarkupKind } from "vscode-languageserver-types";
 import { MarkupContent } from "vscode-languageclient";

apps/vscode/src/providers/format.ts

@@ -44,7 +44,6 @@ import {
   VirtualDoc,
   virtualDocForCode,
   virtualDocForLanguage,
-  virtualDocUri,
   withVirtualDocUri,
 } from "../vdoc/vdoc";
 

apps/vscode/src/vdoc/vdoc-completion.ts

@@ -15,7 +15,7 @@
 
 import { commands, Position, Uri, CompletionList, CompletionItem, Range } from "vscode";
 import { EmbeddedLanguage } from "./languages";
-import { adjustedPosition, unadjustedRange, VirtualDoc, virtualDocUri, withVirtualDocUri } from "./vdoc";
+import { adjustedPosition, unadjustedRange, VirtualDoc, withVirtualDocUri } from "./vdoc";
 
 export async function vdocCompletions(
   vdoc: VirtualDoc,

apps/vscode/src/vdoc/vdoc.ts

@@ -122,14 +122,26 @@ export type VirtualDocAction =
 
 export type VirtualDocUri = { uri: Uri, cleanup?: () => Promise<void> };
 
+/**
+ * Execute a callback on a virtual document's temporary URI
+ *
+ * This method automatically cleans up the temporary URI after executing `f`.
+ *
+ * @param vdoc The virtual document to create a temporary URI for
+ * @param parentUri The virtual document's original URI it was virtualized from
+ * @param f The callback to execute
+ * @returns A Promise evaluating to an object of type `T` returned by `f`
+ */
 export async function withVirtualDocUri<T>(
   vdoc: VirtualDoc,
   parentUri: Uri,
   action: VirtualDocAction,
   f: (uri: Uri) => Promise<T>
-) {
+): Promise<T> {
   const vdocUri = await virtualDocUri(vdoc, parentUri, action);
 
+  // try-finally without a catch allows `f()` to propagate an exception up to the caller
+  // while still allowing us to clean up the vdoc tempfile.
   try {
     return await f(vdocUri.uri);
   } finally {
@@ -139,7 +151,10 @@ export async function withVirtualDocUri<T>(
   }
 }
 
-export async function virtualDocUri(
+// To be used through `withVirtualDocUri()`. Not safe to export on its own! The
+// cleanup hook must be called, and relying on the caller to do this is a huge
+// footgun.
+async function virtualDocUri(
   virtualDoc: VirtualDoc,
   parentUri: Uri,
   action: VirtualDocAction