feat: gate code runs by user whitelist

This commit is contained in:
John
2026-07-02 08:42:44 +08:00
parent 8a53dff856
commit c7d8140a69
7 changed files with 117 additions and 2 deletions
+18
View File
@@ -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) {
+71
View File
@@ -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: {
+2
View File
@@ -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.',
+1
View File
@@ -347,6 +347,7 @@ export function usePageEditSubChat({
resolveAgentRunOptions(trimmed, {
taskType: 'page_edit_code_task',
forceCode: true,
userId: user?.id ?? null,
}),
);
const finishedRun =
+4 -1
View File
@@ -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);
+20 -1
View File
@@ -13,6 +13,23 @@ export const agentCodeRunsAutodetectEnabled = envFlag(
import.meta.env.VITE_AGENT_CODE_RUNS_AUTODETECT,
);
function parseUserIdSet(value: unknown): Set<string> {
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 {};
+1
View File
@@ -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 {