diff --git a/apps/api/src/handlers/tasks/__tests__/updateModelSelection.access.test.ts b/apps/api/src/handlers/tasks/__tests__/updateModelSelection.access.test.ts new file mode 100644 index 000000000..0c1e1c670 --- /dev/null +++ b/apps/api/src/handlers/tasks/__tests__/updateModelSelection.access.test.ts @@ -0,0 +1,251 @@ +import { Hono } from 'hono'; + +import { + automations, + customAutomations, + db, + eq, + inArray, + runFactory, + taskFactory, + taskRuns, + tasks, + userFactory, + users, +} from '@roomote/db/server'; +import { RunStatus } from '@roomote/types'; + +import type { Variables } from '../../../types'; +import type { McpAuth } from '../../mcp/middleware'; + +const { mockWithSandboxServerRpcClient } = vi.hoisted(() => ({ + mockWithSandboxServerRpcClient: vi.fn(), +})); + +vi.mock('@roomote/sdk/server', async (importOriginal) => ({ + ...(await importOriginal()), + withSandboxServerRpcClient: mockWithSandboxServerRpcClient, +})); + +import { updateTaskModelSelection } from '../updateModelSelection'; + +type User = Awaited>; + +const createdAutomationIds: string[] = []; +const createdRunIds: number[] = []; +const createdTaskIds: string[] = []; +const createdUserIds: string[] = []; + +async function createUser(role: 'admin' | 'member') { + const user = await userFactory.create({ role }); + createdUserIds.push(user.id); + return user; +} + +function createApp(user: User) { + const app = new Hono<{ + Variables: Variables & { mcpAuth: McpAuth }; + }>(); + app.use('*', async (c, next) => { + c.set('mcpAuth', { + userId: user.id, + authContext: { + userId: user.id, + tokenType: 'auth', + version: 1, + }, + }); + await next(); + }); + app.post('/tasks/:taskId/model_selection', updateTaskModelSelection); + return app; +} + +function postModelSelection( + app: ReturnType, + taskId: string, + reasoningEffort: 'high' | 'low', +) { + return app.request(`/tasks/${taskId}/model_selection`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + role: 'helper', + model: null, + reasoningEffort, + }), + }); +} + +async function createAutomationTask(owner: User) { + await db + .insert(automations) + .values({ key: 'custom_automation' }) + .onConflictDoNothing(); + const [automation] = await db + .insert(customAutomations) + .values({ + name: `Model selection access ${owner.id}`, + prompt: 'Private task', + createdByUserId: owner.id, + }) + .returning(); + createdAutomationIds.push(automation!.id); + const task = await taskFactory.create({ + initiatorKind: 'automation', + initiatorAutomation: 'custom_automation', + actorExternalId: automation!.id, + }); + createdTaskIds.push(task.id); + const run = await runFactory.create({ + taskId: task.id, + status: RunStatus.Running, + sandboxServerUrl: 'http://sandbox.example.test', + payload: { repo: 'test/repo', description: 'Private task' }, + }); + createdRunIds.push(run.id); + return { task, run }; +} + +async function getRunPayload(runId: number) { + return ( + await db.query.taskRuns.findFirst({ + where: eq(taskRuns.id, runId), + columns: { payload: true }, + }) + )?.payload; +} + +describe('updateTaskModelSelection task access', () => { + beforeEach(() => { + mockWithSandboxServerRpcClient.mockReset(); + mockWithSandboxServerRpcClient.mockResolvedValue({ + application: 'restarted', + }); + }); + + afterEach(async () => { + if (createdRunIds.length > 0) { + await db.delete(taskRuns).where(inArray(taskRuns.id, createdRunIds)); + } + if (createdTaskIds.length > 0) { + await db.delete(tasks).where(inArray(tasks.id, createdTaskIds)); + } + if (createdAutomationIds.length > 0) { + await db + .delete(customAutomations) + .where(inArray(customAutomations.id, createdAutomationIds)); + } + if (createdUserIds.length > 0) { + await db.delete(users).where(inArray(users.id, createdUserIds)); + } + createdRunIds.length = 0; + createdTaskIds.length = 0; + createdAutomationIds.length = 0; + createdUserIds.length = 0; + vi.restoreAllMocks(); + }); + + it('denies another member without persistence or a live restart', async () => { + const owner = await createUser('member'); + const other = await createUser('member'); + const { task, run } = await createAutomationTask(owner); + const originalPayload = await getRunPayload(run.id); + + const response = await postModelSelection( + createApp(other), + task.id, + 'high', + ); + + expect(response.status).toBe(404); + expect(await response.json()).toEqual({ + success: false, + error: 'Task not found', + }); + expect(await getRunPayload(run.id)).toEqual(originalPayload); + expect(mockWithSandboxServerRpcClient).not.toHaveBeenCalled(); + }); + + it.each([ + ['automation owner', 'member'], + ['admin', 'admin'], + ] as const)('allows the %s', async (_label, role) => { + const owner = await createUser('member'); + const actor = role === 'admin' ? await createUser('admin') : owner; + const { task, run } = await createAutomationTask(owner); + + const response = await postModelSelection( + createApp(actor), + task.id, + 'high', + ); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ + success: true, + application: 'restarted', + }); + expect((await getRunPayload(run.id))?.modelRoleOverrides?.helper).toEqual({ + reasoningEffort: 'high', + }); + expect(mockWithSandboxServerRpcClient).toHaveBeenCalledOnce(); + }); + + it('preserves member access to ordinary deployment tasks', async () => { + const owner = await createUser('member'); + const other = await createUser('member'); + const task = await taskFactory.create({ initiatorUserId: owner.id }); + createdTaskIds.push(task.id); + const run = await runFactory.create({ taskId: task.id }); + createdRunIds.push(run.id); + + const response = await postModelSelection( + createApp(other), + task.id, + 'high', + ); + + expect(response.status).toBe(200); + expect((await getRunPayload(run.id))?.modelRoleOverrides?.helper).toEqual({ + reasoningEffort: 'high', + }); + }); + + it('persists a failed live apply and restarts on a later recovery', async () => { + const owner = await createUser('member'); + const { task, run } = await createAutomationTask(owner); + mockWithSandboxServerRpcClient + .mockRejectedValueOnce(new Error('sandbox unavailable')) + .mockResolvedValueOnce({ application: 'restarted' }); + vi.spyOn(console, 'error').mockImplementation(() => {}); + + const failedApply = await postModelSelection( + createApp(owner), + task.id, + 'high', + ); + expect(failedApply.status).toBe(200); + expect(await failedApply.json()).toEqual({ + success: true, + application: 'offline', + }); + expect((await getRunPayload(run.id))?.modelRoleOverrides?.helper).toEqual({ + reasoningEffort: 'high', + }); + + const recoveredApply = await postModelSelection( + createApp(owner), + task.id, + 'low', + ); + expect(await recoveredApply.json()).toEqual({ + success: true, + application: 'restarted', + }); + expect((await getRunPayload(run.id))?.modelRoleOverrides?.helper).toEqual({ + reasoningEffort: 'low', + }); + expect(mockWithSandboxServerRpcClient).toHaveBeenCalledTimes(2); + }); +}); diff --git a/apps/api/src/handlers/tasks/__tests__/updateModelSelection.test.ts b/apps/api/src/handlers/tasks/__tests__/updateModelSelection.test.ts index db004dff1..d0a0a04bc 100644 --- a/apps/api/src/handlers/tasks/__tests__/updateModelSelection.test.ts +++ b/apps/api/src/handlers/tasks/__tests__/updateModelSelection.test.ts @@ -6,10 +6,12 @@ import type { McpAuth } from '../../mcp/middleware'; const { mockApplyTaskModelSelectionToRun, mockFindLatestTaskRun, + mockTaskFindFirst, mockTokenRunFindFirst, } = vi.hoisted(() => ({ mockApplyTaskModelSelectionToRun: vi.fn(), mockFindLatestTaskRun: vi.fn(), + mockTaskFindFirst: vi.fn(), mockTokenRunFindFirst: vi.fn(), })); @@ -28,13 +30,22 @@ vi.mock('@roomote/cloud-agents/server', () => ({ vi.mock('@roomote/db/server', () => ({ db: { query: { + tasks: { + findFirst: mockTaskFindFirst, + }, taskRuns: { findFirst: mockTokenRunFindFirst, }, }, }, + and: vi.fn((...args) => ({ type: 'and', args })), eq: vi.fn((...args) => ({ type: 'eq', args })), taskRuns: { id: 'task_runs.id' }, + tasks: { id: 'tasks.id' }, +})); + +vi.mock('../../custom-automation-history-access', () => ({ + customAutomationHistoryAccess: vi.fn(() => ({ type: 'task-access' })), })); vi.mock('@roomote/sdk/server', () => ({ @@ -83,6 +94,8 @@ describe('updateTaskModelSelection', () => { sandboxServerUrl: null, actingUserId: null, }); + mockTaskFindFirst.mockReset(); + mockTaskFindFirst.mockResolvedValue({ id: 'target-task' }); mockTokenRunFindFirst.mockReset(); }); @@ -134,7 +147,7 @@ describe('updateTaskModelSelection', () => { }); }); - it('applies for a user-token context without a bound run', async () => { + it('applies for an authorized user-token context without a bound run', async () => { const app = createApp({ userId: 'user-1', authContext: { userId: 'user-1' } as never, @@ -144,9 +157,46 @@ describe('updateTaskModelSelection', () => { expect(response.status).toBe(200); expect(mockTokenRunFindFirst).not.toHaveBeenCalled(); + expect(mockTaskFindFirst).toHaveBeenCalledOnce(); expect(mockApplyTaskModelSelectionToRun).toHaveBeenCalled(); }); + it('rejects an unauthorized user without applying', async () => { + mockTaskFindFirst.mockResolvedValue(undefined); + const app = createApp({ + userId: 'user-1', + authContext: { userId: 'user-1' } as never, + }); + + const response = await postModelSelection(app as never, 'target-task'); + + expect(response.status).toBe(404); + expect(await response.json()).toEqual({ + success: false, + error: 'Task not found', + }); + expect(mockFindLatestTaskRun).not.toHaveBeenCalled(); + expect(mockApplyTaskModelSelectionToRun).not.toHaveBeenCalled(); + }); + + it('fails closed when user task authorization cannot be resolved', async () => { + const consoleError = vi + .spyOn(console, 'error') + .mockImplementation(() => {}); + mockTaskFindFirst.mockRejectedValue(new Error('access lookup failed')); + const app = createApp({ + userId: 'user-1', + authContext: { userId: 'user-1' } as never, + }); + + const response = await postModelSelection(app as never, 'target-task'); + + expect(response.status).toBe(500); + expect(mockFindLatestTaskRun).not.toHaveBeenCalled(); + expect(mockApplyTaskModelSelectionToRun).not.toHaveBeenCalled(); + consoleError.mockRestore(); + }); + it('rejects contexts with neither a run token nor a user', async () => { const app = createApp({ userId: undefined, diff --git a/apps/api/src/handlers/tasks/updateModelSelection.ts b/apps/api/src/handlers/tasks/updateModelSelection.ts index 0846df6eb..420e6e955 100644 --- a/apps/api/src/handlers/tasks/updateModelSelection.ts +++ b/apps/api/src/handlers/tasks/updateModelSelection.ts @@ -5,12 +5,13 @@ import { TaskModelSelectionError, applyTaskModelSelectionToRun, } from '@roomote/cloud-agents/server'; -import { db, eq, taskRuns } from '@roomote/db/server'; +import { and, db, eq, taskRuns, tasks } from '@roomote/db/server'; import { withSandboxServerRpcClient } from '@roomote/sdk/server'; import { REASONING_EFFORT_VALUES, isExitedRunStatus } from '@roomote/types'; import type { Variables } from '../../types'; import type { McpAuth } from '../mcp/middleware'; +import { customAutomationHistoryAccess } from '../custom-automation-history-access'; import { isRunTokenContext } from '../mcp/proxy-utils'; import { findLatestTaskRun } from './helpers'; import { logHandlerError } from '../utils'; @@ -41,8 +42,8 @@ const updateModelSelectionBodySchema = z.object({ * A sandbox run token must be bound to the target task: unlike the read and * lifecycle actions on this router, this mutates persistent task state and * force-restarts the harness, so one sandbox must not be able to point it at - * another task (mirrors `manageSourceControl`). User-token contexts carry - * the same authority as the web mutation and pass through. + * another task (mirrors `manageSourceControl`). User-token contexts use the + * same restrictive task-access predicate as sandbox mutations on the web. */ export async function updateTaskModelSelection( c: Context<{ Variables: Variables & { mcpAuth: McpAuth } }>, @@ -76,15 +77,29 @@ export async function updateTaskModelSelection( 403, ); } - } else if (!auth.userId) { - return c.json( - { - success: false, - error: - 'Model selection updates require a task run token or user context', - }, - 403, - ); + } else { + if (!auth.userId) { + return c.json( + { + success: false, + error: + 'Model selection updates require a task run token or user context', + }, + 403, + ); + } + + const task = await db.query.tasks.findFirst({ + where: and( + eq(tasks.id, taskId), + customAutomationHistoryAccess(auth, 'task'), + ), + columns: { id: true }, + }); + + if (!task) { + return c.json({ success: false, error: 'Task not found' }, 404); + } } const parsedBody = updateModelSelectionBodySchema.safeParse(