feat: wire code agent runs to tool mode

This commit is contained in:
John
2026-07-02 07:16:47 +08:00
parent e17dcaf4a8
commit 0271d83756
6 changed files with 343 additions and 22 deletions
+93 -12
View File
@@ -6507,6 +6507,8 @@ async function initSchema(pool2) {
import crypto3 from "node:crypto";
var DEFAULT_RUN_RETRY_DELAYS_MS = [1500, 5e3, 15e3];
var TERMINAL_STATUSES = /* @__PURE__ */ new Set(["succeeded", "failed"]);
var CODE_TOOL_MODES = /* @__PURE__ */ new Set(["code", "code-task", "code_task", "code-tool", "code_tool", "code_tool_task"]);
var RUN_METADATA_KEY = "memindRun";
function nowMs() {
return Date.now();
}
@@ -6520,6 +6522,40 @@ function safeJsonParse(value, fallback = null) {
function serializeMessage(message) {
return JSON.stringify(message ?? {});
}
function normalizeAgentRunToolMode(value) {
const normalized = String(value ?? "chat").trim().toLowerCase();
if (!normalized || normalized === "chat") return "chat";
if (CODE_TOOL_MODES.has(normalized)) return "code";
throw new Error(`\u4E0D\u652F\u6301\u7684 tool_mode: ${value}`);
}
function normalizeTaskType(value) {
const normalized = String(value ?? "").trim();
return normalized || null;
}
function withRunMetadata(userMessage, { toolMode = "chat", taskType = null } = {}) {
const message = userMessage && typeof userMessage === "object" && !Array.isArray(userMessage) ? { ...userMessage } : { value: userMessage };
const metadata = message.metadata && typeof message.metadata === "object" && !Array.isArray(message.metadata) ? { ...message.metadata } : {};
metadata[RUN_METADATA_KEY] = {
...metadata[RUN_METADATA_KEY] && typeof metadata[RUN_METADATA_KEY] === "object" ? metadata[RUN_METADATA_KEY] : {},
toolMode: normalizeAgentRunToolMode(toolMode),
...taskType ? { taskType } : {}
};
return { ...message, metadata };
}
function getRunOptionsFromMessage(userMessage) {
const metadata = userMessage?.metadata;
const runMetadata = metadata?.[RUN_METADATA_KEY] ?? metadata?.agentRun ?? {};
let toolMode = "chat";
try {
toolMode = normalizeAgentRunToolMode(runMetadata?.toolMode ?? metadata?.toolMode ?? "chat");
} catch {
toolMode = "chat";
}
return {
toolMode,
taskType: normalizeTaskType(runMetadata?.taskType ?? metadata?.taskType)
};
}
function projectRun(row) {
if (!row) return null;
return {
@@ -6578,11 +6614,19 @@ function createAgentRunGateway({
);
return rows[0] ?? null;
}
async function createRun(userId, { sessionId = null, requestId, userMessage }) {
async function createRun(userId, {
sessionId = null,
requestId,
userMessage,
toolMode = "chat",
taskType = null
}) {
const normalizedRequestId = String(requestId ?? "").trim();
if (!normalizedRequestId) {
throw new Error("\u7F3A\u5C11 request_id");
}
const normalizedToolMode = normalizeAgentRunToolMode(toolMode);
const normalizedTaskType = normalizeTaskType(taskType);
const existing = await getRunByRequest(userId, normalizedRequestId);
if (existing) {
if (autoDispatch && !TERMINAL_STATUSES.has(existing.status)) dispatchRun(existing.id);
@@ -6590,6 +6634,10 @@ function createAgentRunGateway({
}
const runId = crypto3.randomUUID();
const createdAt = nowMs();
const runMessage = withRunMetadata(userMessage, {
toolMode: normalizedToolMode,
taskType: normalizedTaskType
});
await pool2.query(
`INSERT INTO h5_agent_runs
(id, user_id, agent_session_id, request_id, status, attempts,
@@ -6600,12 +6648,16 @@ function createAgentRunGateway({
userId,
sessionId || null,
normalizedRequestId,
serializeMessage(userMessage),
serializeMessage(runMessage),
createdAt,
createdAt
]
);
await appendEvent(runId, "queued", { sessionId: sessionId || null });
await appendEvent(runId, "queued", {
sessionId: sessionId || null,
toolMode: normalizedToolMode,
taskType: normalizedTaskType
});
if (autoDispatch) dispatchRun(runId);
return projectRun(await getRunById(runId));
}
@@ -6636,22 +6688,32 @@ function createAgentRunGateway({
if (Number(claim?.affectedRows ?? 0) === 0) return;
await appendEvent(runId, "running", { attempt: nextAttempt });
try {
const userMessage = safeJsonParse(row.user_message_json, {});
const runOptions = getRunOptionsFromMessage(userMessage);
let sessionId = row.agent_session_id ?? null;
if (!sessionId) {
const session = await tkmindProxy2.startSessionForUser(row.user_id);
const sessionOptions = {};
if (runOptions.toolMode === "code" && userAuth2?.getCodeAgentSessionPolicy) {
sessionOptions.sessionPolicy = await userAuth2.getCodeAgentSessionPolicy(row.user_id);
}
const session = await tkmindProxy2.startSessionForUser(row.user_id, sessionOptions);
sessionId = session.id;
await pool2.query(
`UPDATE h5_agent_runs SET agent_session_id = ?, updated_at = ? WHERE id = ?`,
[sessionId, nowMs(), runId]
);
await appendEvent(runId, "session_started", { sessionId });
await appendEvent(runId, "session_started", {
sessionId,
toolMode: runOptions.toolMode,
taskType: runOptions.taskType
});
}
const userMessage = safeJsonParse(row.user_message_json, {});
await tkmindProxy2.submitSessionReplyForUser(
row.user_id,
sessionId,
row.request_id,
userMessage
userMessage,
{ toolMode: runOptions.toolMode }
);
await markRun(runId, "succeeded", {
agent_session_id: sessionId,
@@ -6688,6 +6750,8 @@ function createPostAgentRunsHandler({ userAuth: userAuth2, agentRunGateway: agen
const sessionId = String(request.body?.session_id ?? "").trim() || null;
const requestId = String(request.body?.request_id ?? "").trim();
const userMessage = request.body?.user_message ?? null;
const rawToolMode = request.body?.tool_mode ?? request.body?.toolMode ?? "chat";
const taskType = String(request.body?.task_type ?? request.body?.taskType ?? "").trim() || null;
if (!requestId) {
response.status(400).json({ message: "\u7F3A\u5C11 request_id" });
return;
@@ -6696,6 +6760,15 @@ function createPostAgentRunsHandler({ userAuth: userAuth2, agentRunGateway: agen
response.status(400).json({ message: "\u7F3A\u5C11 user_message" });
return;
}
let toolMode = "chat";
try {
toolMode = normalizeAgentRunToolMode(rawToolMode);
} catch (err) {
response.status(400).json({
message: err instanceof Error ? err.message : "\u4E0D\u652F\u6301\u7684 tool_mode"
});
return;
}
if (sessionId) {
const owns = await userAuth2.ownsSession(request.currentUser.id, sessionId);
if (!owns) {
@@ -6706,7 +6779,9 @@ function createPostAgentRunsHandler({ userAuth: userAuth2, agentRunGateway: agen
const run = await agentRunGateway2.createRun(request.currentUser.id, {
sessionId,
requestId,
userMessage
userMessage,
toolMode,
taskType
});
response.status(202).json({ run });
} catch (err) {
@@ -9385,11 +9460,17 @@ function createTkmindProxy({
imgproxySigner
});
}
async function reconcileSessionPolicyForUser(userId, sessionId) {
async function getSessionPolicyForToolMode(userId, toolMode = "chat") {
if (toolMode === "code" && userAuth2.getCodeAgentSessionPolicy) {
return userAuth2.getCodeAgentSessionPolicy(userId);
}
return userAuth2.getAgentSessionPolicy(userId);
}
async function reconcileSessionPolicyForUser(userId, sessionId, { toolMode = "chat" } = {}) {
if (!userId || !sessionId) return;
const target = await resolveTarget(sessionId);
const workingDir = await userAuth2.resolveWorkingDir(userId);
const sessionPolicy = await userAuth2.getAgentSessionPolicy(userId);
const sessionPolicy = await getSessionPolicyForToolMode(userId, toolMode);
const publishLayout = await userAuth2.getUserPublishLayout(userId);
const userMemories = conversationMemoryService2?.listMemories ? await conversationMemoryService2.listMemories(userId, { limit: 40 }).catch(() => []) : [];
await reconcileAgentSession(
@@ -9410,7 +9491,7 @@ function createTkmindProxy({
}
);
}
async function submitSessionReplyForUser(userId, sessionId, requestId, userMessage) {
async function submitSessionReplyForUser(userId, sessionId, requestId, userMessage, { toolMode = "chat" } = {}) {
if (!userId || !sessionId) throw new Error("\u7F3A\u5C11\u4F1A\u8BDD\u4FE1\u606F");
const owns = await userAuth2.ownsSession(userId, sessionId);
if (!owns) throw new Error("\u65E0\u6743\u8BBF\u95EE\u8BE5\u4F1A\u8BDD");
@@ -9421,7 +9502,7 @@ function createTkmindProxy({
err.status = 402;
throw err;
}
await reconcileSessionPolicyForUser(userId, sessionId);
await reconcileSessionPolicyForUser(userId, sessionId, { toolMode });
await applySessionLlmProvider(sessionId);
const user = await userAuth2.getUserById(userId);
if (!user) throw new Error("\u7528\u6237\u4E0D\u5B58\u5728");
+78 -6
View File
@@ -2,6 +2,8 @@ import crypto from 'node:crypto';
const DEFAULT_RUN_RETRY_DELAYS_MS = [1500, 5000, 15000];
const TERMINAL_STATUSES = new Set(['succeeded', 'failed']);
const CODE_TOOL_MODES = new Set(['code', 'code-task', 'code_task', 'code-tool', 'code_tool', 'code_tool_task']);
const RUN_METADATA_KEY = 'memindRun';
function nowMs() {
return Date.now();
@@ -19,6 +21,50 @@ function serializeMessage(message) {
return JSON.stringify(message ?? {});
}
export function normalizeAgentRunToolMode(value) {
const normalized = String(value ?? 'chat').trim().toLowerCase();
if (!normalized || normalized === 'chat') return 'chat';
if (CODE_TOOL_MODES.has(normalized)) return 'code';
throw new Error(`不支持的 tool_mode: ${value}`);
}
function normalizeTaskType(value) {
const normalized = String(value ?? '').trim();
return normalized || null;
}
function withRunMetadata(userMessage, { toolMode = 'chat', taskType = null } = {}) {
const message = (userMessage && typeof userMessage === 'object' && !Array.isArray(userMessage))
? { ...userMessage }
: { value: userMessage };
const metadata = (message.metadata && typeof message.metadata === 'object' && !Array.isArray(message.metadata))
? { ...message.metadata }
: {};
metadata[RUN_METADATA_KEY] = {
...(metadata[RUN_METADATA_KEY] && typeof metadata[RUN_METADATA_KEY] === 'object'
? metadata[RUN_METADATA_KEY]
: {}),
toolMode: normalizeAgentRunToolMode(toolMode),
...(taskType ? { taskType } : {}),
};
return { ...message, metadata };
}
function getRunOptionsFromMessage(userMessage) {
const metadata = userMessage?.metadata;
const runMetadata = metadata?.[RUN_METADATA_KEY] ?? metadata?.agentRun ?? {};
let toolMode = 'chat';
try {
toolMode = normalizeAgentRunToolMode(runMetadata?.toolMode ?? metadata?.toolMode ?? 'chat');
} catch {
toolMode = 'chat';
}
return {
toolMode,
taskType: normalizeTaskType(runMetadata?.taskType ?? metadata?.taskType),
};
}
function projectRun(row) {
if (!row) return null;
return {
@@ -83,11 +129,19 @@ export function createAgentRunGateway({
return rows[0] ?? null;
}
async function createRun(userId, { sessionId = null, requestId, userMessage }) {
async function createRun(userId, {
sessionId = null,
requestId,
userMessage,
toolMode = 'chat',
taskType = null,
}) {
const normalizedRequestId = String(requestId ?? '').trim();
if (!normalizedRequestId) {
throw new Error('缺少 request_id');
}
const normalizedToolMode = normalizeAgentRunToolMode(toolMode);
const normalizedTaskType = normalizeTaskType(taskType);
const existing = await getRunByRequest(userId, normalizedRequestId);
if (existing) {
if (autoDispatch && !TERMINAL_STATUSES.has(existing.status)) dispatchRun(existing.id);
@@ -96,6 +150,10 @@ export function createAgentRunGateway({
const runId = crypto.randomUUID();
const createdAt = nowMs();
const runMessage = withRunMetadata(userMessage, {
toolMode: normalizedToolMode,
taskType: normalizedTaskType,
});
await pool.query(
`INSERT INTO h5_agent_runs
(id, user_id, agent_session_id, request_id, status, attempts,
@@ -106,12 +164,16 @@ export function createAgentRunGateway({
userId,
sessionId || null,
normalizedRequestId,
serializeMessage(userMessage),
serializeMessage(runMessage),
createdAt,
createdAt,
],
);
await appendEvent(runId, 'queued', { sessionId: sessionId || null });
await appendEvent(runId, 'queued', {
sessionId: sessionId || null,
toolMode: normalizedToolMode,
taskType: normalizedTaskType,
});
if (autoDispatch) dispatchRun(runId);
return projectRun(await getRunById(runId));
}
@@ -145,23 +207,33 @@ export function createAgentRunGateway({
await appendEvent(runId, 'running', { attempt: nextAttempt });
try {
const userMessage = safeJsonParse(row.user_message_json, {});
const runOptions = getRunOptionsFromMessage(userMessage);
let sessionId = row.agent_session_id ?? null;
if (!sessionId) {
const session = await tkmindProxy.startSessionForUser(row.user_id);
const sessionOptions = {};
if (runOptions.toolMode === 'code' && userAuth?.getCodeAgentSessionPolicy) {
sessionOptions.sessionPolicy = await userAuth.getCodeAgentSessionPolicy(row.user_id);
}
const session = await tkmindProxy.startSessionForUser(row.user_id, sessionOptions);
sessionId = session.id;
await pool.query(
`UPDATE h5_agent_runs SET agent_session_id = ?, updated_at = ? WHERE id = ?`,
[sessionId, nowMs(), runId],
);
await appendEvent(runId, 'session_started', { sessionId });
await appendEvent(runId, 'session_started', {
sessionId,
toolMode: runOptions.toolMode,
taskType: runOptions.taskType,
});
}
const userMessage = safeJsonParse(row.user_message_json, {});
await tkmindProxy.submitSessionReplyForUser(
row.user_id,
sessionId,
row.request_id,
userMessage,
{ toolMode: runOptions.toolMode },
);
await markRun(runId, 'succeeded', {
agent_session_id: sessionId,
+66
View File
@@ -114,6 +114,34 @@ test('agent run creation is idempotent by user and request id', async () => {
assert.equal(second.status, 'queued');
});
test('agent run creation stores code tool metadata in the queued message', async () => {
const pool = createFakePool();
const gateway = createAgentRunGateway({
pool,
userAuth: {},
tkmindProxy: {},
autoDispatch: false,
});
const run = await gateway.createRun('user-1', {
requestId: 'req-code',
userMessage: {
role: 'user',
content: [{ type: 'text', text: '修改仓库代码' }],
metadata: { source: 'h5' },
},
toolMode: 'code-task',
taskType: 'repo_refactor',
});
const stored = JSON.parse(pool.runs.get(run.id).user_message_json);
assert.equal(stored.metadata.source, 'h5');
assert.deepEqual(stored.metadata.memindRun, {
toolMode: 'code',
taskType: 'repo_refactor',
});
});
test('agent run starts a session and marks submitted reply as succeeded', async () => {
const pool = createFakePool();
const submitted = [];
@@ -141,6 +169,44 @@ test('agent run starts a session and marks submitted reply as succeeded', async
assert.equal(pool.runs.get(run.id).attempts, 1);
});
test('agent run with code tool mode starts and submits with code policy', async () => {
const pool = createFakePool();
const submitted = [];
const codePolicy = {
extensionOverrides: { aider: { allowed: true } },
};
const gateway = createAgentRunGateway({
pool,
userAuth: {
async getCodeAgentSessionPolicy(userId) {
assert.equal(userId, 'user-1');
return codePolicy;
},
},
tkmindProxy: {
async startSessionForUser(userId, options = {}) {
assert.equal(userId, 'user-1');
assert.equal(options.sessionPolicy, codePolicy);
return { id: 'session-code' };
},
async submitSessionReplyForUser(userId, sessionId, requestId, userMessage, options = {}) {
submitted.push({ userId, sessionId, requestId, userMessage, options });
},
},
retryDelaysMs: [],
});
const run = await gateway.createRun('user-1', {
requestId: 'req-code',
userMessage: { role: 'user', content: [{ type: 'text', text: 'run code task' }] },
toolMode: 'code',
});
await waitFor(() => pool.runs.get(run.id)?.status === 'succeeded');
assert.deepEqual(submitted.map((item) => item.options), [{ toolMode: 'code' }]);
assert.equal(submitted[0].userMessage.metadata.memindRun.toolMode, 'code');
});
test('agent run retries transient failures and then becomes terminal', async () => {
const pool = createFakePool();
const gateway = createAgentRunGateway({
+15
View File
@@ -1,9 +1,13 @@
import { normalizeAgentRunToolMode } from './agent-run-gateway.mjs';
export function createPostAgentRunsHandler({ userAuth, agentRunGateway }) {
return async function postAgentRuns(request, response) {
try {
const sessionId = String(request.body?.session_id ?? '').trim() || null;
const requestId = String(request.body?.request_id ?? '').trim();
const userMessage = request.body?.user_message ?? null;
const rawToolMode = request.body?.tool_mode ?? request.body?.toolMode ?? 'chat';
const taskType = String(request.body?.task_type ?? request.body?.taskType ?? '').trim() || null;
if (!requestId) {
response.status(400).json({ message: '缺少 request_id' });
return;
@@ -12,6 +16,15 @@ export function createPostAgentRunsHandler({ userAuth, agentRunGateway }) {
response.status(400).json({ message: '缺少 user_message' });
return;
}
let toolMode = 'chat';
try {
toolMode = normalizeAgentRunToolMode(rawToolMode);
} catch (err) {
response.status(400).json({
message: err instanceof Error ? err.message : '不支持的 tool_mode',
});
return;
}
if (sessionId) {
const owns = await userAuth.ownsSession(request.currentUser.id, sessionId);
if (!owns) {
@@ -23,6 +36,8 @@ export function createPostAgentRunsHandler({ userAuth, agentRunGateway }) {
sessionId,
requestId,
userMessage,
toolMode,
taskType,
});
response.status(202).json({ run });
} catch (err) {
+74
View File
@@ -93,11 +93,85 @@ test('POST /agent/runs creates a run and returns 202', async () => {
sessionId: 'session-1',
requestId: 'req-1',
userMessage: { role: 'user', content: [] },
toolMode: 'chat',
taskType: null,
},
},
]);
});
test('POST /agent/runs accepts explicit code tool mode and task type', 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' };
},
},
});
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-task',
task_type: 'repo_refactor',
},
},
res,
);
assert.equal(res.statusCode, 202);
assert.deepEqual(created[0].payload, {
sessionId: 'session-1',
requestId: 'req-code',
userMessage: { role: 'user', content: [] },
toolMode: 'code',
taskType: 'repo_refactor',
});
});
test('POST /agent/runs rejects unsupported tool mode', async () => {
const handler = createPostAgentRunsHandler({
userAuth: {
async ownsSession() {
throw new Error('should not check ownership after invalid mode');
},
},
agentRunGateway: {
async createRun() {
throw new Error('should not be called');
},
},
});
const res = createResponseRecorder();
await handler(
{
currentUser: { id: 'user-1' },
body: {
request_id: 'req-1',
user_message: { role: 'user', content: [] },
tool_mode: 'admin',
},
},
res,
);
assert.equal(res.statusCode, 400);
assert.match(res.body.message, /tool_mode/);
});
test('POST /agent/runs rejects a session the user does not own', async () => {
const handler = createPostAgentRunsHandler({
userAuth: {
+17 -4
View File
@@ -1061,11 +1061,18 @@ export function createTkmindProxy({
});
}
async function reconcileSessionPolicyForUser(userId, sessionId) {
async function getSessionPolicyForToolMode(userId, toolMode = 'chat') {
if (toolMode === 'code' && userAuth.getCodeAgentSessionPolicy) {
return userAuth.getCodeAgentSessionPolicy(userId);
}
return userAuth.getAgentSessionPolicy(userId);
}
async function reconcileSessionPolicyForUser(userId, sessionId, { toolMode = 'chat' } = {}) {
if (!userId || !sessionId) return;
const target = await resolveTarget(sessionId);
const workingDir = await userAuth.resolveWorkingDir(userId);
const sessionPolicy = await userAuth.getAgentSessionPolicy(userId);
const sessionPolicy = await getSessionPolicyForToolMode(userId, toolMode);
const publishLayout = await userAuth.getUserPublishLayout(userId);
const userMemories = conversationMemoryService?.listMemories
? await conversationMemoryService.listMemories(userId, { limit: 40 }).catch(() => [])
@@ -1091,7 +1098,13 @@ export function createTkmindProxy({
);
}
async function submitSessionReplyForUser(userId, sessionId, requestId, userMessage) {
async function submitSessionReplyForUser(
userId,
sessionId,
requestId,
userMessage,
{ toolMode = 'chat' } = {},
) {
if (!userId || !sessionId) throw new Error('缺少会话信息');
const owns = await userAuth.ownsSession(userId, sessionId);
if (!owns) throw new Error('无权访问该会话');
@@ -1103,7 +1116,7 @@ export function createTkmindProxy({
throw err;
}
await reconcileSessionPolicyForUser(userId, sessionId);
await reconcileSessionPolicyForUser(userId, sessionId, { toolMode });
await applySessionLlmProvider(sessionId);
const user = await userAuth.getUserById(userId);