diff --git a/agent-run-routes.mjs b/agent-run-routes.mjs index 5649a46..f23941d 100644 --- a/agent-run-routes.mjs +++ b/agent-run-routes.mjs @@ -4,10 +4,20 @@ function envFlag(value) { return ['1', 'true', 'yes', 'on'].includes(String(value ?? '').trim().toLowerCase()); } +function parseUserIdSet(value) { + return new Set( + String(value ?? '') + .split(',') + .map((item) => item.trim()) + .filter(Boolean), + ); +} + export function createPostAgentRunsHandler({ userAuth, agentRunGateway, codeRunsEnabled = envFlag(process.env.MEMIND_AGENT_CODE_RUNS_ENABLED), + codeRunUserIds = parseUserIdSet(process.env.MEMIND_AGENT_CODE_RUNS_USER_IDS), }) { return async function postAgentRuns(request, response) { try { @@ -37,6 +47,14 @@ export function createPostAgentRunsHandler({ response.status(403).json({ message: '代码任务灰度未开启' }); return; } + if ( + toolMode === 'code' && + codeRunUserIds.size > 0 && + !codeRunUserIds.has(request.currentUser.id) + ) { + response.status(403).json({ message: '当前用户未开启代码任务灰度' }); + return; + } if (sessionId) { const owns = await userAuth.ownsSession(request.currentUser.id, sessionId); if (!owns) { diff --git a/agent-run-routes.test.mjs b/agent-run-routes.test.mjs index d8c40b2..ec1af37 100644 --- a/agent-run-routes.test.mjs +++ b/agent-run-routes.test.mjs @@ -142,6 +142,77 @@ test('POST /agent/runs accepts explicit code tool mode and task type', async () }); }); +test('POST /agent/runs accepts code tool mode for whitelisted user', async () => { + const created = []; + const handler = createPostAgentRunsHandler({ + userAuth: { + async ownsSession() { + return true; + }, + }, + agentRunGateway: { + async createRun(userId, payload) { + created.push({ userId, payload }); + return { id: 'run-code', status: 'queued' }; + }, + }, + codeRunsEnabled: true, + codeRunUserIds: new Set(['user-1']), + }); + const res = createResponseRecorder(); + + await handler( + { + currentUser: { id: 'user-1' }, + body: { + session_id: 'session-1', + request_id: 'req-code', + user_message: { role: 'user', content: [] }, + tool_mode: 'code', + }, + }, + res, + ); + + assert.equal(res.statusCode, 202); + assert.equal(created[0].userId, 'user-1'); + assert.equal(created[0].payload.toolMode, 'code'); +}); + +test('POST /agent/runs rejects code tool mode for non-whitelisted user', async () => { + const handler = createPostAgentRunsHandler({ + userAuth: { + async ownsSession() { + throw new Error('should not check ownership after user gate rejects code mode'); + }, + }, + agentRunGateway: { + async createRun() { + throw new Error('should not be called'); + }, + }, + codeRunsEnabled: true, + codeRunUserIds: new Set(['user-2']), + }); + const res = createResponseRecorder(); + + await handler( + { + currentUser: { id: 'user-1' }, + body: { + session_id: 'session-1', + request_id: 'req-code', + user_message: { role: 'user', content: [] }, + tool_mode: 'code', + }, + }, + res, + ); + + assert.equal(res.statusCode, 403); + assert.deepEqual(res.body, { message: '当前用户未开启代码任务灰度' }); +}); + test('POST /agent/runs rejects code tool mode when server gate is disabled', async () => { const handler = createPostAgentRunsHandler({ userAuth: { diff --git a/scripts/check-agent-code-run-entry.mjs b/scripts/check-agent-code-run-entry.mjs index a4ec79c..df47876 100644 --- a/scripts/check-agent-code-run-entry.mjs +++ b/scripts/check-agent-code-run-entry.mjs @@ -144,6 +144,8 @@ const buildFlags = { serverCodeRunsEnabledTruthy: truthy(process.env.MEMIND_AGENT_CODE_RUNS_ENABLED), enabledEnv: process.env.VITE_AGENT_CODE_RUNS_ENABLED ?? null, autodetectEnv: process.env.VITE_AGENT_CODE_RUNS_AUTODETECT ?? null, + enabledUserIdsEnv: process.env.VITE_AGENT_CODE_RUNS_USER_IDS ?? null, + serverCodeRunUserIdsEnv: process.env.MEMIND_AGENT_CODE_RUNS_USER_IDS ?? null, enabledTruthy: truthy(process.env.VITE_AGENT_CODE_RUNS_ENABLED), autodetectTruthy: truthy(process.env.VITE_AGENT_CODE_RUNS_AUTODETECT), note: 'These env values describe the current shell, not necessarily the already-built dist.', diff --git a/src/hooks/usePageEditSubChat.ts b/src/hooks/usePageEditSubChat.ts index 9f25975..c315ea0 100644 --- a/src/hooks/usePageEditSubChat.ts +++ b/src/hooks/usePageEditSubChat.ts @@ -347,6 +347,7 @@ export function usePageEditSubChat({ resolveAgentRunOptions(trimmed, { taskType: 'page_edit_code_task', forceCode: true, + userId: user?.id ?? null, }), ); const finishedRun = diff --git a/src/hooks/useTKMindChat.ts b/src/hooks/useTKMindChat.ts index 8acf5ce..4f07bc7 100644 --- a/src/hooks/useTKMindChat.ts +++ b/src/hooks/useTKMindChat.ts @@ -1173,7 +1173,10 @@ export function useTKMindChat( activeSessionId, requestId, userMessage, - resolveAgentRunOptions(trimmed, { taskType: 'h5_chat_code_task' }), + resolveAgentRunOptions(trimmed, { + taskType: 'h5_chat_code_task', + userId: userRef.current?.id ?? null, + }), ); const finishedRun = createdRun.status === 'succeeded' ? createdRun : await waitForAgentRun(createdRun.id); diff --git a/src/utils/agentRunMode.ts b/src/utils/agentRunMode.ts index ad16718..819824b 100644 --- a/src/utils/agentRunMode.ts +++ b/src/utils/agentRunMode.ts @@ -13,6 +13,23 @@ export const agentCodeRunsAutodetectEnabled = envFlag( import.meta.env.VITE_AGENT_CODE_RUNS_AUTODETECT, ); +function parseUserIdSet(value: unknown): Set { + return new Set( + String(value ?? '') + .split(',') + .map((item) => item.trim()) + .filter(Boolean), + ); +} + +const agentCodeRunUserIds = parseUserIdSet(import.meta.env.VITE_AGENT_CODE_RUNS_USER_IDS); + +export function agentCodeRunsEnabledForUser(userId?: string | null): boolean { + if (!agentCodeRunsEnabled) return false; + if (agentCodeRunUserIds.size === 0) return true; + return Boolean(userId && agentCodeRunUserIds.has(userId)); +} + const CODE_TASK_PATTERNS = [ /\b(repo|repository|branch|commit|pull request|pr|diff|patch)\b/i, /\b(aider|openhands|codex|codebase|workspace)\b/i, @@ -27,13 +44,15 @@ export function resolveAgentRunOptions( taskType = 'code_task', forceCode = false, allowAutodetect = agentCodeRunsAutodetectEnabled, + userId = null, }: { taskType?: string; forceCode?: boolean; allowAutodetect?: boolean; + userId?: string | null; } = {}, ): AgentRunCreateOptions { - if (!agentCodeRunsEnabled) return {}; + if (!agentCodeRunsEnabledForUser(userId)) return {}; const normalizedText = String(text ?? '').trim(); const shouldUseCode = forceCode || (allowAutodetect && CODE_TASK_PATTERNS.some((pattern) => pattern.test(normalizedText))); if (!shouldUseCode) return {}; diff --git a/src/vite-env.d.ts b/src/vite-env.d.ts index 1eeca6e..ad50b3e 100644 --- a/src/vite-env.d.ts +++ b/src/vite-env.d.ts @@ -9,6 +9,7 @@ interface ImportMetaEnv { readonly VITE_MINDSPACE_BASE?: string; readonly VITE_AGENT_CODE_RUNS_ENABLED?: string; readonly VITE_AGENT_CODE_RUNS_AUTODETECT?: string; + readonly VITE_AGENT_CODE_RUNS_USER_IDS?: string; } interface ImportMeta {