improvement(computer-use): roll out computer use per organization (#9182)

* improvement(computer-use): roll out computer use per organization

* test(computer-use): drop assertions that restate mocked rollout calls

Author: waleedlatif1

apps/sim/app/api/desktop/computer/availability/route.ts

@@ -12,10 +12,10 @@ export const GET = defineInternalJsonRoute({
   auth: internalSessionAuth,
   operation: readComputerUseAvailability.operation,
   rateLimit: internalRateLimits.none({
-    reason: 'Read-only rollout switch used on each desktop chat admission.',
+    reason: 'Read-only rollout check used on each desktop chat admission.',
   }),
   errorPolicy: internalOrchestrationErrorPolicy,
-  mapInput: () => undefined,
+  mapInput: ({ query }) => query,
   useCase: readComputerUseAvailability,
   present: (result) => result,
 })

apps/sim/app/o/[organizationId]/layout.tsx

@@ -88,7 +88,7 @@ export default async function OrganizationLayout({
       logger.warn('Could not resolve benchmark navigation access')
       return false
     }),
-    isComputerUseAvailable(),
+    isComputerUseAvailable(organizationId),
     isWorkflowTestsEnabled(organizationId),
     isChangelogEnabled(organizationId),
     isProjectsEnabled(),

apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts

@@ -3929,7 +3929,16 @@ export function useChat(
         const desktopChatCapabilities =
           options?.requestMode === 'assistant'
             ? {}
-            : await getDesktopChatCapabilities(desktopScopeIdRef.current, computerUseEnabled)
+            : await getDesktopChatCapabilities(
+                desktopScopeIdRef.current,
+                !computerUseEnabled
+                  ? null
+                  : organizationId
+                    ? { organizationId }
+                    : workspaceId
+                      ? { workspaceId }
+                      : null
+              )
 
         admissionUnknown = true
         /** A reload from here on may find the server holding this id. */

apps/sim/app/workspace/[workspaceId]/layout.tsx

@@ -150,7 +150,7 @@ async function WorkspaceContent({
     isWorkflowTestsEnabled(hostContext.hostOrganizationId),
     isChangelogEnabled(hostContext.hostOrganizationId),
     hasDesktopBackgroundExecutor(session.user.id),
-    isComputerUseAvailable(),
+    isComputerUseAvailable(hostContext.hostOrganizationId),
     isProjectsEnabled(),
     prefetchWorkspaceAccess(queryClient, workspaceId, principal),
     prefetchWorkspaceForkAvailability(queryClient, workspaceId, principal, hostContext),

apps/sim/lib/api/contracts/computer-use.ts

@@ -1,14 +1,22 @@
 import { ComputerUseSchema } from '@sim/desktop-bridge/computer-use'
 import { z } from 'zod'
 import { desktopToolCallIdSchema } from '@/lib/api/contracts/desktop-tool-authorization'
-import { defineRouteContract } from '@/lib/api/contracts/types'
+import { organizationIdSchema, workspaceIdSchema } from '@/lib/api/contracts/primitives'
+import { type ContractQuery, defineRouteContract } from '@/lib/api/contracts/types'
 
+/** The chat owner whose rollout applies: exactly one of a workspace or an organization. */
+const computerUseAvailabilityQuerySchema = z.union([
+  z.object({ workspaceId: workspaceIdSchema }).strict(),
+  z.object({ organizationId: organizationIdSchema }).strict(),
+])
 const computerUseAvailabilityResponseSchema = z.object({ enabled: z.boolean() }).strict()
 export const computerUseAvailabilityContract = defineRouteContract({
   method: 'GET',
   path: '/api/desktop/computer/availability',
+  query: computerUseAvailabilityQuerySchema,
   response: { mode: 'json', schema: computerUseAvailabilityResponseSchema },
 })
+export type ComputerUseOwner = ContractQuery<typeof computerUseAvailabilityContract>
 
 const authorizeComputerUseBodySchema = z.object({ toolCallId: desktopToolCallIdSchema }).strict()
 const authorizeComputerUseResponseSchema = z

apps/sim/lib/computer-use/application/authorize.ts

@@ -60,7 +60,7 @@ export const authorizeComputerUse = defineAuthorizedChatUseCase({
   },
   authorizationOptions: {},
   async execute({ principal, context, input }) {
-    if (!(await isComputerUseAvailable()))
+    if (!(await isComputerUseAvailable(context.organizationId ?? context.workspaceOrganizationId)))
       throw new OrchestrationError('forbidden', 'Computer use is unavailable on this deployment')
     const claimed = await claimComputerUseTool({
       toolCallId: input.toolCallId,

apps/sim/lib/computer-use/application/availability.ts

@@ -1,26 +1,77 @@
+import type { Principal } from '@sim/auth/principal'
+import type { ComputerUseOwner } from '@/lib/api/contracts/computer-use'
 import { isComputerUseAvailable } from '@/lib/computer-use/availability.server'
+import { defineAuthorizedWorkspaceUseCase, defineWorkspaceOperation } from '@/lib/core/application'
+import { defineAuthorizedOrganizationUseCase } from '@/lib/core/application/authorized-organization-use-case'
 import {
   assertOperationPrincipal,
   defineOperation,
   type OperationUseCase,
 } from '@/lib/core/application/operation'
+import { defineOrganizationOperation } from '@/lib/core/application/organization-operation'
+import { resolveActiveWorkspaceApplicationContext } from '@/lib/workspaces/application/context'
 
-/** permission-group-exempt: reports only the deployment-wide rollout switch; actions require chat authorization. */
+/** permission-group-exempt: reports the rollout for a workspace the caller can read; actions require chat authorization. */
+const workspaceAvailabilityOperation = defineWorkspaceOperation({
+  id: 'desktop.computer.availability',
+  minimumRole: 'read',
+  workspaceApiKey: 'deny',
+  capability: 'none',
+  principalKinds: ['session'],
+})
+
+/** permission-group-exempt: reports the rollout for an organization the caller belongs to; actions require chat authorization. */
+const organizationAvailabilityOperation = defineOrganizationOperation({
+  id: 'desktop.computer.availability',
+  minimumRole: 'member',
+  capability: 'none',
+  principalKinds: ['session'],
+})
+
+const workspaceAvailability = defineAuthorizedWorkspaceUseCase({
+  operation: workspaceAvailabilityOperation,
+  resolveContext: ({ input }: { input: { workspaceId: string } }) =>
+    resolveActiveWorkspaceApplicationContext(input.workspaceId),
+  authorizationOptions: {},
+  execute: ({ context }) => isComputerUseAvailable(context.workspaceOrganizationId),
+})
+
+const organizationAvailability = defineAuthorizedOrganizationUseCase({
+  operation: organizationAvailabilityOperation,
+  execute: ({ input }: { input: { organizationId: string } }) =>
+    isComputerUseAvailable(input.organizationId),
+})
+
+/**
+ * Whether computer use is rolled out for a chat's owner: an organization chat follows its
+ * organization, a workspace chat its workspace's organization, and a personal workspace the
+ * global switch.
+ */
+export function isComputerUseAvailableFor(
+  principal: Principal,
+  owner: ComputerUseOwner
+): Promise<boolean> {
+  return 'organizationId' in owner
+    ? organizationAvailability.execute({ principal, input: owner })
+    : workspaceAvailability.execute({ principal, input: owner })
+}
+
+/** permission-group-exempt: dispatches to the owner's authorized rollout read above. */
 const availabilityOperation = defineOperation({
   id: 'desktop.computer.availability',
   capability: 'none',
   principalKinds: ['session'],
 })
 
-/** The authenticated settings and capability read; no device or app data leaves the desktop. */
+/** The desktop's rollout read before it touches the native helper; no device or app data leaves the desktop. */
 export const readComputerUseAvailability: OperationUseCase<
   typeof availabilityOperation,
-  undefined,
+  ComputerUseOwner,
   { enabled: boolean }
 > = {
   operation: availabilityOperation,
-  async execute({ principal }) {
+  async execute({ principal, input }) {
     assertOperationPrincipal(principal, availabilityOperation)
-    return { enabled: await isComputerUseAvailable() }
+    return { enabled: await isComputerUseAvailableFor(principal, input) }
   },
 }

apps/sim/lib/computer-use/availability.server.ts

@@ -1,6 +1,6 @@
 import { isFeatureEnabled } from '@/lib/core/config/feature-flags'
 
-/** Re-evaluates the global rollout at every server admission boundary. */
-export async function isComputerUseAvailable(): Promise<boolean> {
-  return isFeatureEnabled('mothership-computer-use')
+/** Computer use rolls out by organization; personal workspaces follow the global switch. */
+export function isComputerUseAvailable(orgId?: string | null): Promise<boolean> {
+  return isFeatureEnabled('mothership-computer-use', orgId ? { orgId } : {})
 }

apps/sim/lib/core/config/feature-flags.ts

@@ -53,7 +53,8 @@ const FEATURE_FLAGS = {
   },
   'mothership-computer-use': {
     description:
-      'Enable native macOS computer use in Mothership. Global on/off only; each device must also opt in.',
+      'Enable native macOS computer use in Mothership. Supports global and organization rollout; ' +
+      'personal workspaces follow the global switch. Each device must also opt in.',
     fallback: 'MSHIP_COMPUTER_USE',
   },
   'mothership-search-integration-tools': {

apps/sim/lib/desktop/index.test.ts

@@ -62,7 +62,7 @@ describe('desktop surface availability', () => {
         terminalEnabled: false,
       })
       requestAvailability.mockResolvedValueOnce({ enabled: rollout })
-      const result = await getDesktopChatCapabilities('chat-1', true)
+      const result = await getDesktopChatCapabilities('chat-1', { workspaceId: 'ws-1' })
       expect(result.desktopCapabilities?.computerUse ?? false).toBe(expected)
     }
   )
@@ -75,7 +75,7 @@ describe('desktop surface availability', () => {
       browserAgent: { getKnownSessions: vi.fn(async () => ({ sessions: [] })) },
     })
     requestAvailability.mockResolvedValue({ enabled: false })
-    const result = await getDesktopChatCapabilities('chat-1', false)
+    const result = await getDesktopChatCapabilities('chat-1', null)
     expect(requestAvailability).not.toHaveBeenCalled()
     expect(getStatus).not.toHaveBeenCalled()
     expect(result.desktopCapabilities).toMatchObject({
@@ -95,7 +95,7 @@ describe('desktop surface availability', () => {
       terminalEnabled: false,
     })
     requestAvailability.mockReturnValueOnce(rollout.promise)
-    const result = getDesktopChatCapabilities('chat-1', true)
+    const result = getDesktopChatCapabilities('chat-1', { workspaceId: 'ws-1' })
     await vi.waitFor(() => expect(requestAvailability).toHaveBeenCalledOnce())
     expect(getStatus).not.toHaveBeenCalled()
     rollout.resolve({ enabled: false })
@@ -113,7 +113,8 @@ describe('desktop surface availability', () => {
     })
     requestAvailability.mockRejectedValueOnce(new Error('offline'))
     expect(
-      (await getDesktopChatCapabilities('chat-1', true)).desktopCapabilities?.computerUse
+      (await getDesktopChatCapabilities('chat-1', { workspaceId: 'ws-1' })).desktopCapabilities
+        ?.computerUse
     ).toBeUndefined()
     expect(getStatus).not.toHaveBeenCalled()
   })
@@ -143,12 +144,12 @@ describe('desktop surface availability', () => {
       browserEnabled: false,
       terminalEnabled: false,
     })
-    expect(await getDesktopChatCapabilities('org-chat', false)).toMatchObject({
+    expect(await getDesktopChatCapabilities('org-chat', null)).toMatchObject({
       desktopCapabilities: { localFiles: true },
     })
     installBridge({})
     expect(
-      (await getDesktopChatCapabilities('org-chat', false)).desktopCapabilities?.localFiles
+      (await getDesktopChatCapabilities('org-chat', null)).desktopCapabilities?.localFiles
     ).toBeUndefined()
   })
 
@@ -160,7 +161,9 @@ describe('desktop surface availability', () => {
     })
     const offerFrom = async (bridge: Record<string, unknown>) => {
       installBridge({ localFiles: vi.fn(), ...bridge })
-      const { desktopCapabilities } = await getDesktopChatCapabilities('chat-1', true)
+      const { desktopCapabilities } = await getDesktopChatCapabilities('chat-1', {
+        workspaceId: 'ws-1',
+      })
       return { deviceId: desktopCapabilities?.deviceId, executor: desktopCapabilities?.executor }
     }
     const unset = { deviceId: undefined, executor: undefined }
@@ -208,7 +211,7 @@ describe('desktop surface availability', () => {
       browserEnabled: false,
     })
 
-    const capabilities = await getDesktopChatCapabilities('chat-1', true)
+    const capabilities = await getDesktopChatCapabilities('chat-1', { workspaceId: 'ws-1' })
 
     expect(capabilities.desktopCapabilities?.terminals).toEqual([
       {