fix: resolve project context for untitled query files (#1827)

Untitled files created via 'New query' or manually silently failed to
execute or compile because the selected project was never persisted to
workspace state, and executeSQL/compileQuery could not resolve untitled
URIs to a dbt project.

Changes:
- createSqlFile now stores the selected project in workspace state
- Extract resolveProjectUri() helper with safe Uri reconstruction from
  deserialized workspace state; apply to both executeSQL and compileQuery
- Add ensureProjectForUntitledUri() fallback in RunModel that validates
  stored project still exists (clears stale state) and falls back to
  getOrPickProjectFromWorkspace for untitled files not created via
  createSqlFile
- Extract toProjectQuickPickItem() to eliminate duplicated shape across
  3 call sites
- Invalidate dbtPowerUser.projectSelected when a project is unregistered
- Use clean 'untitled' modelName instead of raw fsPath for untitled URIs
- Wire QueryManifestService into RunModel via DI

Fixes #1827

Author: shreyastelkar

src/commands/index.ts

@@ -33,7 +33,10 @@ import { DBTProject } from "../dbt_client/dbtProject";
 import { DBTProjectContainer } from "../dbt_client/dbtProjectContainer";
 import { PythonEnvironment } from "../dbt_client/pythonEnvironment";
 import { NotebookQuickPick } from "../quickpick/notebookQuickPick";
-import { ProjectQuickPickItem } from "../quickpick/projectQuickPick";
+import {
+  ProjectQuickPickItem,
+  toProjectQuickPickItem,
+} from "../quickpick/projectQuickPick";
 import { DiagnosticsOutputChannel } from "../services/diagnosticsOutputChannel";
 import { QueryManifestService } from "../services/queryManifestService";
 import { RunHistoryService } from "../services/runHistoryService";
@@ -697,6 +700,13 @@ export class VSCodeCommands implements Disposable {
               return;
             }
 
+            // Persist project selection so untitled files can resolve
+            // project context during query execution and compilation
+            this.dbtProjectContainer.setToWorkspaceState(
+              "dbtPowerUser.projectSelected",
+              toProjectQuickPickItem(project),
+            );
+
             // Open a new untitled sql file by default
             let docOpenPromise = workspace.openTextDocument({
               language: "jinja-sql",

src/commands/runModel.ts

@@ -3,11 +3,16 @@ import { RunModelType } from "@altimateai/dbt-integration";
 import { Uri, window } from "vscode";
 import { GenerateModelFromSourceParams } from "../code_lens_provider/sourceModelCreationCodeLensProvider";
 import { DBTProjectContainer } from "../dbt_client/dbtProjectContainer";
+import { toProjectQuickPickItem } from "../quickpick/projectQuickPick";
+import { QueryManifestService } from "../services/queryManifestService";
 import { NodeTreeItem } from "../treeview_provider/modelTreeviewProvider";
 import { extendErrorWithSupportLinks } from "../utils";
 
 export class RunModel {
-  constructor(private dbtProjectContainer: DBTProjectContainer) {}
+  constructor(
+    private dbtProjectContainer: DBTProjectContainer,
+    private queryManifestService: QueryManifestService,
+  ) {}
 
   runModelOnActiveWindow(type?: RunModelType) {
     if (!window.activeTextEditor) {
@@ -41,17 +46,61 @@ export class RunModel {
     this.compileDBTModel(fullPath);
   }
 
-  compileQueryOnActiveWindow() {
+  async compileQueryOnActiveWindow() {
     if (!window.activeTextEditor) {
       return;
     }
     const fullPath = window.activeTextEditor.document.uri;
+    if (fullPath.scheme === "untitled") {
+      const resolved = await this.ensureProjectForUntitledUri();
+      if (!resolved) {
+        return;
+      }
+    }
     const query = window.activeTextEditor.document.getText();
     if (query !== undefined) {
       this.compileDBTQuery(fullPath, query);
     }
   }
 
+  /**
+   * Ensures a dbt project is stored in workspace state for untitled files.
+   * If no project is stored yet — or the stored project is stale (removed
+   * between sessions) — falls back to getOrPickProjectFromWorkspace() which
+   * auto-selects the single project or prompts the user to pick one.
+   * Returns true if a project is available, false if the user cancelled or
+   * no projects exist.
+   */
+  private async ensureProjectForUntitledUri(): Promise<boolean> {
+    const selectedProject = this.dbtProjectContainer.getFromWorkspaceState(
+      "dbtPowerUser.projectSelected",
+    );
+    if (selectedProject?.uri) {
+      // Validate the stored project still exists in the workspace
+      const raw = selectedProject.uri;
+      const storedUri = Uri.file(raw.fsPath || raw.path);
+      if (this.dbtProjectContainer.findDBTProject(storedUri)) {
+        return true;
+      }
+      // Stale — clear it and fall through to re-pick
+      this.dbtProjectContainer.setToWorkspaceState(
+        "dbtPowerUser.projectSelected",
+        undefined,
+      );
+    }
+    const project =
+      await this.queryManifestService.getOrPickProjectFromWorkspace();
+    if (!project) {
+      window.showErrorMessage("Unable to find dbt project for this query.");
+      return false;
+    }
+    this.dbtProjectContainer.setToWorkspaceState(
+      "dbtPowerUser.projectSelected",
+      toProjectQuickPickItem(project),
+    );
+    return true;
+  }
+
   private getQuery() {
     if (!window.activeTextEditor) {
       return;
@@ -62,16 +111,26 @@ export class RunModel {
     );
   }
 
-  executeQueryOnActiveWindow() {
+  async executeQueryOnActiveWindow() {
     const query = this.getQuery();
     if (query === undefined) {
       return;
     }
     const modelPath = window.activeTextEditor?.document.uri;
-    if (modelPath) {
-      const modelName = path.basename(modelPath.fsPath, ".sql");
-      this.executeSQL(window.activeTextEditor!.document.uri, query, modelName);
+    if (!modelPath) {
+      return;
+    }
+    if (modelPath.scheme === "untitled") {
+      const resolved = await this.ensureProjectForUntitledUri();
+      if (!resolved) {
+        return;
+      }
     }
+    const modelName =
+      modelPath.scheme === "untitled"
+        ? "untitled"
+        : path.basename(modelPath.fsPath, ".sql");
+    this.executeSQL(modelPath, query, modelName);
   }
 
   runModelOnNodeTreeItem(type: RunModelType) {

src/dbt_client/dbtProjectContainer.ts

@@ -102,6 +102,7 @@ export class DBTProjectContainer implements Disposable {
         this.projects.set(event.root, event.name);
       } else {
         this.projects.delete(event.root);
+        this.clearStaleProjectSelection(event.root);
       }
     });
   }
@@ -240,15 +241,10 @@ export class DBTProjectContainer implements Disposable {
   }
 
   executeSQL(uri: Uri, query: string, modelName: string): void {
-    if (uri.scheme === "untitled") {
-      const selectedProject = this.getFromWorkspaceState(
-        "dbtPowerUser.projectSelected",
-      );
-      if (selectedProject) {
-        uri = selectedProject.uri;
-      }
-    }
-    this.findDBTProject(uri)?.executeSQLOnQueryPanel(query, modelName);
+    this.findDBTProject(this.resolveProjectUri(uri))?.executeSQLOnQueryPanel(
+      query,
+      modelName,
+    );
   }
 
   runModel(modelPath: Uri, type?: RunModelType) {
@@ -286,7 +282,9 @@ export class DBTProjectContainer implements Disposable {
   }
 
   compileQuery(modelPath: Uri, query: string) {
-    return this.findDBTProject(modelPath)?.compileQuery(query);
+    return this.findDBTProject(this.resolveProjectUri(modelPath))?.compileQuery(
+      query,
+    );
   }
 
   showRunSQL(modelPath: Uri) {
@@ -479,6 +477,49 @@ export class DBTProjectContainer implements Disposable {
     return this.dbtWorkspaceFolders.find((folder) => folder.contains(uri));
   }
 
+  /**
+   * Resolves an untitled document URI to the selected project's URI.
+   * When users create ad-hoc query files (via "New query" or manually),
+   * the document has an `untitled:` scheme that can't be matched to a
+   * project directory. This method looks up the stored project selection
+   * from workspace state and returns a real file URI that findDBTProject
+   * can resolve.
+   *
+   * Workspace state stores Uri as a plain JSON object (not a Uri instance),
+   * so we reconstruct it via Uri.file() to ensure .fsPath works correctly.
+   */
+  private resolveProjectUri(uri: Uri): Uri {
+    if (uri.scheme === "untitled") {
+      const selectedProject = this.getFromWorkspaceState(
+        "dbtPowerUser.projectSelected",
+      );
+      if (selectedProject?.uri) {
+        const raw = selectedProject.uri;
+        return Uri.file(raw.fsPath || raw.path);
+      }
+    }
+    return uri;
+  }
+
+  /**
+   * Clear the stored project selection if it points to a project that was
+   * just unregistered. Prevents stale state from causing silent failures
+   * when the user later opens an untitled query file.
+   */
+  private clearStaleProjectSelection(removedRoot: Uri): void {
+    const selectedProject = this.getFromWorkspaceState(
+      "dbtPowerUser.projectSelected",
+    );
+    if (!selectedProject?.uri) {
+      return;
+    }
+    const raw = selectedProject.uri;
+    const selectedPath = raw.fsPath || raw.path;
+    if (selectedPath === removedRoot.fsPath) {
+      this.setToWorkspaceState("dbtPowerUser.projectSelected", undefined);
+    }
+  }
+
   async checkIfAltimateDatapilotInstalled() {
     const datapilotVersion =
       await this.altimateDatapilot.checkIfAltimateDatapilotInstalled();

src/inversify.config.ts

@@ -1415,7 +1415,10 @@ container
 container
   .bind(RunModel)
   .toDynamicValue((context) => {
-    return new RunModel(context.container.get(DBTProjectContainer));
+    return new RunModel(
+      context.container.get(DBTProjectContainer),
+      context.container.get(QueryManifestService),
+    );
   })
   .inSingletonScope();
 

src/quickpick/projectQuickPick.ts

@@ -7,17 +7,24 @@ export interface ProjectQuickPickItem extends QuickPickItem {
   uri: Uri;
 }
 
+/** Build a ProjectQuickPickItem from a DBTProject for workspace state storage. */
+export function toProjectQuickPickItem(
+  project: DBTProject,
+): ProjectQuickPickItem {
+  return {
+    label: project.getProjectName(),
+    description: project.projectRoot.fsPath,
+    uri: project.projectRoot,
+  };
+}
+
 export class ProjectQuickPick {
   async projectPicker(
     projects: DBTProject[],
   ): Promise<ProjectQuickPickItem | undefined> {
-    const options: ProjectQuickPickItem[] = projects.map((item) => {
-      return {
-        label: item.getProjectName(),
-        description: item.projectRoot.fsPath,
-        uri: item.projectRoot,
-      };
-    });
+    const options: ProjectQuickPickItem[] = projects.map(
+      toProjectQuickPickItem,
+    );
 
     const pick: ProjectQuickPickItem | undefined = await window.showQuickPick(
       options,