fix(workspace): guard prompt billing context

Co-authored-by: Amp <amp@ampcode.com>
Amp-Thread-ID: https://ampcode.com/threads/T-01a01a57-6259-75bd-bb52-c149e0f7bb9c

Author: dante01yoon

src/locales/en/main.json

@@ -2128,7 +2128,8 @@
     "extensionFileHint": "This may be due to the following script",
     "promptExecutionError": "Prompt execution failed",
     "accessRestrictedTitle": "Access Restricted",
-    "accessRestrictedMessage": "Your account is not authorized for this feature."
+    "accessRestrictedMessage": "Your account is not authorized for this feature.",
+    "workspaceChangedDuringExecution": "Your active workspace changed while preparing this execution. Run the workflow again to use the current workspace."
   },
   "apiNodesSignInDialog": {
     "title": "Sign In Required to Use API Nodes",

src/scripts/app.test.ts

@@ -60,6 +60,7 @@ const {
   mockToastStore,
   mockExtensionService,
   mockNodeOutputStore,
+  mockTeamWorkspaceStore,
   mockWorkspaceWorkflow,
   mockRefreshMissingModelPipeline,
   mockImportA1111,
@@ -87,6 +88,9 @@ const {
     refreshNodeOutputs: vi.fn(),
     resetAllOutputsAndPreviews: vi.fn()
   },
+  mockTeamWorkspaceStore: {
+    activeWorkspaceId: 'workspace-a' as string | null
+  },
   mockWorkspaceWorkflow: {
     activeWorkflow: null as ComfyWorkflow | null,
     createNewTemporary: vi.fn(),
@@ -119,6 +123,10 @@ vi.mock('@/stores/authStore', () => ({
   useAuthStore: vi.fn(() => mockAuthStore)
 }))
 
+vi.mock('@/platform/workspace/stores/teamWorkspaceStore', () => ({
+  useTeamWorkspaceStore: vi.fn(() => mockTeamWorkspaceStore)
+}))
+
 vi.mock('@/platform/settings/settingStore', () => ({
   useSettingStore: vi.fn(() => mockSettingStore)
 }))
@@ -281,6 +289,7 @@ describe('ComfyApp', () => {
     })
     mockApiKeyAuthStore.getApiKey.mockReturnValue(undefined)
     mockAuthStore.getWorkspaceAuthToken.mockResolvedValue(undefined)
+    mockTeamWorkspaceStore.activeWorkspaceId = 'workspace-a'
     mockExtensionService.invokeExtensions.mockReturnValue([])
     mockExtensionService.invokeExtensionsAsync.mockResolvedValue(undefined)
     vi.mocked(extractFilesFromDragEvent).mockResolvedValue([])
@@ -335,6 +344,66 @@ describe('ComfyApp', () => {
       expect(queuePrompt).toHaveBeenCalledOnce()
     })
 
+    it('uses a workspace initialized while authentication is pending', async () => {
+      prepareEmptyPromptQueue()
+      mockTeamWorkspaceStore.activeWorkspaceId = null
+      mockAuthStore.getWorkspaceAuthToken.mockImplementationOnce(async () => {
+        mockTeamWorkspaceStore.activeWorkspaceId = 'workspace-a'
+        return 'workspace-token'
+      })
+      const queuePrompt = vi
+        .spyOn(api, 'queuePrompt')
+        .mockImplementation(() => {
+          expect(api.authToken).toBe('workspace-token')
+          return Promise.resolve({ prompt_id: 'job-1', error: '' })
+        })
+
+      await expect(app.queuePrompt(0)).resolves.toBe(true)
+
+      expect(queuePrompt).toHaveBeenCalledOnce()
+    })
+
+    it('does not submit when the workspace changes during authentication', async () => {
+      prepareEmptyPromptQueue()
+      mockAuthStore.getWorkspaceAuthToken.mockImplementationOnce(async () => {
+        mockTeamWorkspaceStore.activeWorkspaceId = 'workspace-b'
+        return 'workspace-token-a'
+      })
+      const queuePrompt = vi.spyOn(api, 'queuePrompt')
+      const showDialog = vi.spyOn(useDialogStore(), 'showDialog')
+
+      await expect(app.queuePrompt(0)).resolves.toBe(false)
+
+      expect(queuePrompt).not.toHaveBeenCalled()
+      expect(showDialog).toHaveBeenCalledOnce()
+    })
+
+    it('does not submit a prompt after the active workspace changes', async () => {
+      prepareEmptyPromptQueue()
+      let finishPromptBuild: () => void = () => {}
+      vi.spyOn(app, 'graphToPrompt').mockImplementationOnce(
+        () =>
+          new Promise((resolve) => {
+            finishPromptBuild = () =>
+              resolve({
+                output: {},
+                workflow: createWorkflowGraphData()
+              })
+          })
+      )
+      const queuePrompt = vi.spyOn(api, 'queuePrompt')
+      const showDialog = vi.spyOn(useDialogStore(), 'showDialog')
+
+      const submission = app.queuePrompt(0)
+      await vi.waitFor(() => expect(app.graphToPrompt).toHaveBeenCalledOnce())
+      mockTeamWorkspaceStore.activeWorkspaceId = 'workspace-b'
+      finishPromptBuild()
+
+      await expect(submission).resolves.toBe(false)
+      expect(queuePrompt).not.toHaveBeenCalled()
+      expect(showDialog).toHaveBeenCalledOnce()
+    })
+
     it('preserves missing node packs when submitting a prompt', async () => {
       prepareEmptyPromptQueue()
       const missingNodesStore = useMissingNodesErrorStore()

src/scripts/app.ts

@@ -79,6 +79,7 @@ import {
 } from '@/scripts/domWidget'
 import { useAccountPreconditionDialog } from '@/platform/cloud/subscription/composables/useAccountPreconditionDialog'
 import { resolveAccountPrecondition } from '@/platform/errorCatalog/accountPreconditionRouting'
+import { useTeamWorkspaceStore } from '@/platform/workspace/stores/teamWorkspaceStore'
 import { useDialogService } from '@/services/dialogService'
 import { useExtensionService } from '@/services/extensionService'
 import { useLitegraphService } from '@/services/litegraphService'
@@ -1640,7 +1641,13 @@ export class ComfyApp {
     let queueResultOverride: boolean | null = null
 
     // Get auth token for backend nodes - uses workspace token if enabled, otherwise Firebase token
+    const workspaceIdBeforeAuthentication =
+      useTeamWorkspaceStore().activeWorkspaceId
     const comfyOrgAuthToken = await useAuthStore().getWorkspaceAuthToken()
+    const executionWorkspaceId = useTeamWorkspaceStore().activeWorkspaceId
+    const workspaceChangedWhileAuthenticating =
+      workspaceIdBeforeAuthentication !== null &&
+      workspaceIdBeforeAuthentication !== executionWorkspaceId
     const comfyOrgApiKey = useApiKeyAuthStore().getApiKey()
 
     try {
@@ -1713,6 +1720,20 @@ export class ComfyApp {
               viewMode: getWorkflowMode(queuedWorkflow)
             })
           }
+          if (
+            workspaceChangedWhileAuthenticating ||
+            executionWorkspaceId !== useTeamWorkspaceStore().activeWorkspaceId
+          ) {
+            useDialogService().showErrorDialog(
+              new Error(t('errorDialog.workspaceChangedDuringExecution')),
+              {
+                title: t('errorDialog.promptExecutionError'),
+                reportType: 'promptExecutionError'
+              }
+            )
+            queueResultOverride = false
+            break
+          }
           try {
             api.authToken = comfyOrgAuthToken
             api.apiKey = comfyOrgApiKey ?? undefined