diff --git a/apps/api/src/handlers/tasks/__tests__/launchTask.test.ts b/apps/api/src/handlers/tasks/__tests__/launchTask.test.ts index 2d25759ee..7ad9277df 100644 --- a/apps/api/src/handlers/tasks/__tests__/launchTask.test.ts +++ b/apps/api/src/handlers/tasks/__tests__/launchTask.test.ts @@ -10,6 +10,7 @@ const { mockEnqueueTask, mockEnvironmentsFindFirst, mockRepositoriesFindMany, + mockTaskRunsFindFirst, mockSelectRows, mockResolveWorkspaceRepositoryProviders, mockGetMembershipRole, @@ -17,6 +18,7 @@ const { mockEnqueueTask: vi.fn(), mockEnvironmentsFindFirst: vi.fn(), mockRepositoriesFindMany: vi.fn(), + mockTaskRunsFindFirst: vi.fn(), mockSelectRows: vi.fn(), mockResolveWorkspaceRepositoryProviders: vi.fn(), mockGetMembershipRole: vi.fn(), @@ -37,6 +39,7 @@ vi.mock('@roomote/db/server', () => ({ environments: {}, environmentRepositoryMappings: {}, repositories: {}, + taskRuns: {}, resolveWorkspaceRepositoryProviders: (...args: unknown[]) => mockResolveWorkspaceRepositoryProviders(...args), db: { @@ -47,6 +50,9 @@ vi.mock('@roomote/db/server', () => ({ repositories: { findMany: (...args: unknown[]) => mockRepositoriesFindMany(...args), }, + taskRuns: { + findFirst: (...args: unknown[]) => mockTaskRunsFindFirst(...args), + }, }, select: () => { const chain = { @@ -90,6 +96,8 @@ describe('launchTask', () => { mockEnvironmentsFindFirst.mockReset(); mockEnvironmentsFindFirst.mockResolvedValue({ id: 'env-1' }); mockRepositoriesFindMany.mockReset(); + mockTaskRunsFindFirst.mockReset(); + mockTaskRunsFindFirst.mockResolvedValue(undefined); mockSelectRows.mockReset(); mockSelectRows.mockReturnValue([]); mockResolveWorkspaceRepositoryProviders.mockReset(); @@ -299,6 +307,68 @@ describe('launchTask', () => { expect(enqueuedTask.task.payload.notifySourceRunOnSettle).toBe(true); }); + it.each(['docker', 'modal'] as const)( + 'inherits the %s source run compute provider for run-token child launches', + async (provider) => { + mockEnqueueTask.mockResolvedValue({ id: 103, taskId: 'task-child' }); + mockTaskRunsFindFirst.mockResolvedValue({ vendor: provider }); + + const runAuth = { + runId: 555, + userId: 'user-1', + principal: 'user', + tokenType: 'run', + version: 1, + } as RunTokenContext; + + const app = createApp(runAuth); + const response = await app.request( + new Request('http://localhost/tasks', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ prompt: 'Verify the environment' }), + }), + ); + + expect(response.status).toBe(200); + const enqueuedTask = mockEnqueueTask.mock.calls[0]?.[0] as { + task: { computeProvider?: string }; + }; + expect(enqueuedTask.task.computeProvider).toBe(provider); + }, + ); + + it('preserves an explicit compute provider on run-token child launches', async () => { + mockEnqueueTask.mockResolvedValue({ id: 104, taskId: 'task-child' }); + + const runAuth = { + runId: 556, + userId: 'user-1', + principal: 'user', + tokenType: 'run', + version: 1, + } as RunTokenContext; + + const app = createApp(runAuth); + const response = await app.request( + new Request('http://localhost/tasks', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + prompt: 'Verify the environment', + computeProvider: 'modal', + }), + }), + ); + + expect(response.status).toBe(200); + const enqueuedTask = mockEnqueueTask.mock.calls[0]?.[0] as { + task: { computeProvider?: string }; + }; + expect(enqueuedTask.task.computeProvider).toBe('modal'); + expect(mockTaskRunsFindFirst).not.toHaveBeenCalled(); + }); + it('ignores notifyOnSettle for user-token launches', async () => { mockEnqueueTask.mockResolvedValue({ id: 103, taskId: 'task-user' }); diff --git a/apps/api/src/handlers/tasks/launchTask.ts b/apps/api/src/handlers/tasks/launchTask.ts index 94fd91fd2..0d60ed5e7 100644 --- a/apps/api/src/handlers/tasks/launchTask.ts +++ b/apps/api/src/handlers/tasks/launchTask.ts @@ -13,11 +13,13 @@ import { inArray, repositories, resolveWorkspaceRepositoryProviders, + taskRuns, } from '@roomote/db/server'; import { ADMIN_REQUIRED_LAUNCH_TYPES, ALL_REPOSITORIES, buildTaskTypePromptAndWorkspacePayload, + type ComputeProvider, getEnvironmentRepositoryInstallationError, type StandardTask, type SuggestedTasksTask, @@ -125,6 +127,25 @@ async function resolveLaunchSourceControlProvider({ return undefined; } +async function resolveLaunchComputeProvider({ + requestedProvider, + auth, +}: { + requestedProvider: ComputeProvider | undefined; + auth: McpAuth; +}): Promise { + if (requestedProvider || !('runId' in auth.authContext)) { + return requestedProvider; + } + + const sourceRun = await db.query.taskRuns.findFirst({ + where: eq(taskRuns.id, auth.authContext.runId), + columns: { vendor: true }, + }); + + return sourceRun?.vendor ?? undefined; +} + /** * POST /api/tasks * @@ -284,6 +305,10 @@ export async function launchTask( requestedType === 'standard' ? body.bootstrap?.skill : undefined, userId: auth.userId, }); + const computeProvider = await resolveLaunchComputeProvider({ + requestedProvider: body.computeProvider, + auth, + }); // A settle notification needs a durable pointer back to the launching // run, so the opt-in only takes effect on run-token launches. @@ -294,7 +319,7 @@ export async function launchTask( const taskBase = { harness: harnessSelection.harness ?? body.harness, - computeProvider: body.computeProvider, + computeProvider, requestedWorkKindDecision, ...((requestedType === 'environment-definition' || notifySourceRunOnSettle) &&