feat(auth): add workspace env override

Author: AmanVarshney01

docs/product/command-spec.md

@@ -63,16 +63,19 @@ Out of scope for the current beta:
 
 ## Authentication
 
-The CLI accepts two authentication sources, in this fixed precedence:
+The CLI accepts these authentication inputs, in this fixed precedence:
 
 1. `PRISMA_SERVICE_TOKEN` environment variable — long-lived service token, intended for CI and other headless contexts.
-2. Stored OAuth session — created by `prisma-cli auth login`, kept in the OS-appropriate credentials store, refreshed automatically.
+2. `PRISMA_CLI_WORKSPACE_ID` environment variable — process-local selector for one locally authenticated OAuth workspace.
+3. Stored active OAuth workspace — created by `prisma-cli auth login`, kept in the OS-appropriate credentials store, refreshed automatically.
 
 Stored OAuth sessions include a short-lived access token and a refresh token. The local credentials store may contain OAuth grants for multiple workspaces. One local active workspace pointer selects which grant authenticated commands use. Commands refresh the selected access token automatically when the API rejects it, coordinate refreshes across concurrent CLI processes, and tolerate short refresh-token rotation races. If the selected session cannot be refreshed, commands fail with a structured `AUTH_REQUIRED` error instead of surfacing SDK stack traces or silently falling through to another workspace.
 
 When `PRISMA_SERVICE_TOKEN` is set and non-empty, the token is fully sufficient for authenticated commands. If `PRISMA_SERVICE_TOKEN` is set but empty or only whitespace, commands fail with an auth configuration error instead of falling back to stored OAuth. The CLI does not read any locally stored OAuth session when a non-empty service token is present, so behavior is identical on a fresh runner and a developer machine that happens to be signed in. The active workspace is derived from the token's `sub` claim; no additional flag or environment variable is required for the common case where the token is scoped to a single workspace.
 
-`auth login`, `auth logout`, and `auth workspace` operate on stored OAuth sessions. They do not affect the `PRISMA_SERVICE_TOKEN` environment variable. `auth login` stores the authorized workspace and makes it active. `auth logout` clears all local OAuth workspace sessions. `auth workspace logout` and `auth logout --workspace` clear one local OAuth workspace session, including while `PRISMA_SERVICE_TOKEN` is set, because they only clean local OAuth state. If that workspace was active, the CLI does not silently fall through to another cached workspace; the user must explicitly choose the next workspace with `auth workspace use`. `auth workspace use` changes only local CLI context and never mutates a remote resource. When `PRISMA_SERVICE_TOKEN` is set, workspace switching is unavailable because the token is the active auth source.
+When `PRISMA_CLI_WORKSPACE_ID` is set and `PRISMA_SERVICE_TOKEN` is not set, authenticated commands use the matching locally stored OAuth workspace for that process only. The value matches a workspace id or canonical workspace id from `auth workspace list`, including the same id with or without a `wksp_` prefix; it does not match workspace names. This environment variable does not mutate the stored active workspace pointer, which makes it suitable for parallel agents or scripts that should not fight over shared CLI state. If it is empty, ambiguous, or does not match a locally authenticated workspace, commands fail with a structured auth error instead of falling back to another workspace.
+
+`auth login`, `auth logout`, and `auth workspace` operate on stored OAuth sessions. They do not affect the `PRISMA_SERVICE_TOKEN` environment variable. `auth login` stores the authorized workspace and makes it active. `auth logout` clears all local OAuth workspace sessions. `auth workspace logout` and `auth logout --workspace` clear one local OAuth workspace session, including while `PRISMA_SERVICE_TOKEN` is set, because they only clean local OAuth state. If that workspace was active, the CLI does not silently fall through to another cached workspace; the user must explicitly choose the next workspace with `auth workspace use`. `auth workspace use` changes only local CLI context and never mutates a remote resource; its output also includes a `PRISMA_CLI_WORKSPACE_ID=<id> prisma-cli project list` example for process-local use. When `PRISMA_SERVICE_TOKEN` is set, workspace switching is unavailable because the token is the active auth source.
 
 ## Context Resolution
 

packages/cli/src/adapters/token-storage.ts

@@ -4,7 +4,11 @@ import fs from "node:fs/promises";
 import path from "node:path";
 import { CredentialsStore } from "@prisma/credentials-store";
 import type { TokenStorage, Tokens } from "@prisma/management-api-sdk";
-import { getAuthFilePath } from "../lib/auth/client";
+import {
+  getAuthFilePath,
+  getWorkspaceIdOverride,
+  WORKSPACE_ID_ENV_VAR,
+} from "../lib/auth/client";
 
 interface StoredCredential {
   workspaceId?: unknown;
@@ -26,6 +30,10 @@ export interface StoredAuthWorkspaceLogout {
   activeWorkspace: StoredAuthWorkspace | null;
 }
 
+interface ListWorkspacesOptions {
+  migrateAuthContext?: boolean;
+}
+
 interface AuthContextWorkspace {
   id?: unknown;
   name?: unknown;
@@ -162,7 +170,7 @@ export class FileTokenStorage implements TokenStorage {
   private readonly lockFilePath: string;
 
   constructor(
-    env: NodeJS.ProcessEnv = process.env,
+    private readonly env: NodeJS.ProcessEnv = process.env,
     private readonly signal?: AbortSignal,
     private readonly options: FileTokenStorageOptions = {},
   ) {
@@ -176,8 +184,18 @@ export class FileTokenStorage implements TokenStorage {
   async getTokens(): Promise<Tokens | null> {
     this.signal?.throwIfAborted();
     try {
+      const workspaceIdOverride = getWorkspaceIdOverride(this.env);
       // CredentialsStore does not accept AbortSignal; check immediately before and after the boundary.
       const credentials = await this.readCredentialsFromDisk();
+
+      if (workspaceIdOverride) {
+        const context = await this.readAuthContext();
+        return this.selectWorkspaceTokens(credentials, context, {
+          ref: workspaceIdOverride,
+          matcher: workspaceMatchesIdRef,
+        });
+      }
+
       const context = await this.readAuthContext();
 
       if (context.state.activeWorkspaceId) {
@@ -199,8 +217,9 @@ export class FileTokenStorage implements TokenStorage {
       // state, usually after logging out the active workspace. Do not fall back
       // to another cached workspace without an explicit `auth workspace use`.
       return null;
-    } catch (_error) {
+    } catch (error) {
       if (this.signal?.aborted) throw this.signal.reason;
+      if (isWorkspaceOverrideError(error)) throw error;
       return null;
     }
   }
@@ -268,38 +287,17 @@ export class FileTokenStorage implements TokenStorage {
     this.signal?.throwIfAborted();
   }
 
-  async listWorkspaces(): Promise<StoredAuthWorkspace[]> {
+  async listWorkspaces(
+    options: ListWorkspacesOptions = {},
+  ): Promise<StoredAuthWorkspace[]> {
     this.signal?.throwIfAborted();
     const credentials = await this.readCredentialsFromDisk();
-    const context = await this.ensureMigratedAuthContext(credentials);
+    const context =
+      options.migrateAuthContext === false
+        ? await this.readAuthContext()
+        : await this.ensureMigratedAuthContext(credentials);
 
-    return credentials
-      .map((credential) => storedCredentialToTokens(credential))
-      .filter((tokens): tokens is Tokens => tokens !== null)
-      .map((tokens) => {
-        const cached = context.state.workspaces[tokens.workspaceId];
-        const id =
-          typeof cached?.id === "string" && cached.id.trim().length > 0
-            ? cached.id.trim()
-            : tokens.workspaceId;
-        const name =
-          typeof cached?.name === "string" && cached.name.trim().length > 0
-            ? workspaceDisplayName(cached.name.trim(), tokens.workspaceId)
-            : UNKNOWN_WORKSPACE_NAME;
-        const lastSeenAt =
-          typeof cached?.lastSeenAt === "string" &&
-          cached.lastSeenAt.trim().length > 0
-            ? cached.lastSeenAt.trim()
-            : null;
-
-        return {
-          id,
-          name,
-          credentialWorkspaceId: tokens.workspaceId,
-          active: context.state.activeWorkspaceId === tokens.workspaceId,
-          lastSeenAt,
-        };
-      });
+    return this.buildStoredAuthWorkspaces(credentials, context);
   }
 
   async listWorkspaceTokens(): Promise<Tokens[]> {
@@ -310,6 +308,21 @@ export class FileTokenStorage implements TokenStorage {
       .filter((tokens): tokens is Tokens => tokens !== null);
   }
 
+  async getTokensForWorkspaceId(workspaceId: string): Promise<Tokens | null> {
+    this.signal?.throwIfAborted();
+    const ref = workspaceId.trim();
+    if (!ref) {
+      throw new WorkspaceSelectionError("missing");
+    }
+
+    const credentials = await this.readCredentialsFromDisk();
+    const context = await this.readAuthContext();
+    return this.selectWorkspaceTokens(credentials, context, {
+      ref,
+      matcher: workspaceMatchesIdRef,
+    });
+  }
+
   async useWorkspace(workspaceRef: string): Promise<{
     previous: StoredAuthWorkspace | null;
     selected: StoredAuthWorkspace;
@@ -431,6 +444,14 @@ export class FileTokenStorage implements TokenStorage {
     workspace: { id: string; name: string },
   ): Promise<void> {
     const context = await this.readAuthContext();
+    if (
+      !context.exists &&
+      this.env[WORKSPACE_ID_ENV_VAR] !== undefined &&
+      this.options.activateOnSetTokens !== true
+    ) {
+      return;
+    }
+
     context.state.workspaces[credentialWorkspaceId] = {
       id: workspace.id,
       name: workspace.name,
@@ -546,6 +567,66 @@ export class FileTokenStorage implements TokenStorage {
     return (await this.credentialsStore.getCredentials()) as StoredCredential[];
   }
 
+  private selectWorkspaceTokens(
+    credentials: StoredCredential[],
+    context: AuthContextReadResult,
+    options: {
+      ref: string;
+      matcher: (workspace: StoredAuthWorkspace, ref: string) => boolean;
+    },
+  ): Tokens | null {
+    const workspaces = this.buildStoredAuthWorkspaces(credentials, context);
+    const matches = workspaces.filter((workspace) =>
+      options.matcher(workspace, options.ref),
+    );
+
+    if (matches.length === 0) {
+      throw new WorkspaceSelectionError("not-found", options.ref);
+    }
+
+    if (matches.length > 1) {
+      throw new WorkspaceSelectionError("ambiguous", options.ref, matches);
+    }
+
+    return findTokensForWorkspace(
+      credentials,
+      matches[0].credentialWorkspaceId,
+    );
+  }
+
+  private buildStoredAuthWorkspaces(
+    credentials: StoredCredential[],
+    context: AuthContextReadResult,
+  ): StoredAuthWorkspace[] {
+    return credentials
+      .map((credential) => storedCredentialToTokens(credential))
+      .filter((tokens): tokens is Tokens => tokens !== null)
+      .map((tokens) => {
+        const cached = context.state.workspaces[tokens.workspaceId];
+        const id =
+          typeof cached?.id === "string" && cached.id.trim().length > 0
+            ? cached.id.trim()
+            : tokens.workspaceId;
+        const name =
+          typeof cached?.name === "string" && cached.name.trim().length > 0
+            ? workspaceDisplayName(cached.name.trim(), tokens.workspaceId)
+            : UNKNOWN_WORKSPACE_NAME;
+        const lastSeenAt =
+          typeof cached?.lastSeenAt === "string" &&
+          cached.lastSeenAt.trim().length > 0
+            ? cached.lastSeenAt.trim()
+            : null;
+
+        return {
+          id,
+          name,
+          credentialWorkspaceId: tokens.workspaceId,
+          active: context.state.activeWorkspaceId === tokens.workspaceId,
+          lastSeenAt,
+        };
+      });
+  }
+
   private async ensureMigratedAuthContext(
     credentials: StoredCredential[],
   ): Promise<AuthContextReadResult> {
@@ -584,8 +665,6 @@ export class FileTokenStorage implements TokenStorage {
     const context = await this.readAuthContext();
     if (
       this.options.activateOnSetTokens === false &&
-      context.exists &&
-      context.state.activeWorkspaceId &&
       context.state.activeWorkspaceId !== workspaceId
     ) {
       return;
@@ -605,6 +684,14 @@ export class FileTokenStorage implements TokenStorage {
     options: { preserveActivePointer: boolean },
   ): Promise<void> {
     const context = await this.readAuthContext();
+    if (
+      !context.exists &&
+      this.env[WORKSPACE_ID_ENV_VAR] !== undefined &&
+      options.preserveActivePointer
+    ) {
+      return;
+    }
+
     delete context.state.workspaces[workspaceId];
     if (
       !options.preserveActivePointer &&
@@ -711,6 +798,30 @@ function workspaceMatchesRef(
   );
 }
 
+export function workspaceMatchesIdRef(
+  workspace: StoredAuthWorkspace,
+  ref: string,
+): boolean {
+  return (
+    workspace.credentialWorkspaceId === ref ||
+    workspace.id === ref ||
+    stripWorkspacePrefix(workspace.credentialWorkspaceId) ===
+      stripWorkspacePrefix(ref) ||
+    stripWorkspacePrefix(workspace.id) === stripWorkspacePrefix(ref)
+  );
+}
+
+function isWorkspaceOverrideError(error: unknown): boolean {
+  if (error instanceof WorkspaceSelectionError) {
+    return true;
+  }
+
+  return (
+    error instanceof Error &&
+    error.message.startsWith(`${WORKSPACE_ID_ENV_VAR} is set but empty`)
+  );
+}
+
 function stripWorkspacePrefix(value: string): string {
   return value.startsWith("wksp_") ? value.slice("wksp_".length) : value;
 }