From 24adde79fefd31d4efc332a5c774c847b151377a Mon Sep 17 00:00:00 2001 From: Deng Junhui <120795714+minorcell@users.noreply.github.com> Date: Tue, 16 Dec 2025 03:15:46 +0800 Subject: [PATCH] feat: add hook middleware support --- packages/core/README.md | 11 +++ packages/core/src/runtime/hooks.test.ts | 111 ++++++++++++++++++++++++ packages/core/src/runtime/session.ts | 76 +++++++++++++++- packages/core/src/types.ts | 50 ++++++++++- packages/ui/src/index.ts | 2 +- 5 files changed, 246 insertions(+), 4 deletions(-) create mode 100644 packages/core/src/runtime/hooks.test.ts diff --git a/packages/core/README.md b/packages/core/README.md index dcae852c..6b07764a 100644 --- a/packages/core/README.md +++ b/packages/core/README.md @@ -36,3 +36,14 @@ await session.close() ``` 如果提供自定义工具/LLM/prompt/sink，可在 deps/options 中覆盖对应字段。默认配置会选择当前 provider 并写入用户目录的 sessions。 + +## Hook 与中间件 + +Core 支持在关键阶段注册 Hook，并可通过 middleware 链式扩展： + +- `onTurnStart(payload, next?)`：每轮开始时触发，`payload` 包含 `turn` 与用户输入。 +- `onAction(payload, next?)`：模型请求调用工具时触发，可用于审计参数或落盘。 +- `onObservation(payload, next?)`：工具返回后触发，`payload` 提供 `tool`、`observation`、`step` 等信息。 +- `onFinal(payload, next?)`：回合结束时触发，附带最终文本与状态码。 + +Hook 既可以直接传入回调，也可以传入 `(next, payload) => { ...; await next() }` 形式的中间件数组，满足链路追踪、埋点等需求。所有 Hook 均为可选，未提供时不会影响主流程。 diff --git a/packages/core/src/runtime/hooks.test.ts b/packages/core/src/runtime/hooks.test.ts new file mode 100644 index 00000000..895c80ee --- /dev/null +++ b/packages/core/src/runtime/hooks.test.ts @@ -0,0 +1,111 @@ +import { beforeAll, afterAll, describe, expect, test } from 'bun:test' +import { join } from 'node:path' +import { tmpdir } from 'node:os' +import { $ } from 'bun' +import { z } from 'zod' +import { createAgentSession, createTokenCounter } from '@memo/core' +import type { AgentSession } from '@memo/core' + +let tempHome: string +let previousHome: string | undefined + +async function makeTempDir(prefix: string) { + const dir = join(tmpdir(), `${prefix}-${crypto.randomUUID()}`) + await $`mkdir -p ${dir}` + return dir +} + +async function removeDir(dir: string) { + await $`rm -rf ${dir}` +} + +describe('hook middleware chain', () => { + let session: AgentSession | undefined + + beforeAll(async () => { + tempHome = await makeTempDir('memo-core-hooks') + previousHome = process.env.MEMO_HOME + process.env.MEMO_HOME = tempHome + }) + + afterAll(async () => { + if (session) { + await session.close() + } + if (previousHome === undefined) { + delete process.env.MEMO_HOME + } else { + process.env.MEMO_HOME = previousHome + } + await removeDir(tempHome) + }) + + test('runs optional hooks in order with middleware', async () => { + const calls: string[] = [] + let llmStep = 0 + + session = await createAgentSession( + { + tools: { + echo: { + name: 'echo', + description: 'echo text', + inputSchema: z.object({ text: z.string() }), + async execute(input) { + return { content: [{ type: 'text', text: `echo:${input.text}` }] } as any + }, + }, + }, + callLLM: async () => { + llmStep += 1 + if (llmStep === 1) { + return JSON.stringify({ tool: 'echo', input: { text: 'hello' } }) + } + return JSON.stringify({ final: 'done' }) + }, + historySinks: [], + tokenCounter: createTokenCounter('cl100k_base'), + onTurnStart: [ + ({ turn, input }) => { + calls.push(`start:${turn}:${input}`) + }, + async (next, payload) => { + calls.push('start-mw-before') + await next() + calls.push('start-mw-after') + }, + ], + onAction: ({ tool, input }) => { + calls.push(`action:${tool}:${JSON.stringify(input)}`) + }, + onObservation: [ + ({ observation }) => { + calls.push(`ob:${observation}`) + }, + async (next) => { + calls.push('ob-mw-before') + await next() + calls.push('ob-mw-after') + }, + ], + onFinal: ({ final, status }) => { + calls.push(`final:${final}:${status}`) + }, + }, + { mode: 'once' }, + ) + + const result = await session.runTurn('hi') + expect(result.finalText).toBe('done') + expect(calls).toEqual([ + 'start:1:hi', + 'start-mw-before', + 'start-mw-after', + 'action:echo:{"text":"hello"}', + 'ob:echo:hello', + 'ob-mw-before', + 'ob-mw-after', + 'final:done:ok', + ]) + }) +}) diff --git a/packages/core/src/runtime/session.ts b/packages/core/src/runtime/session.ts index 2a57435a..defa0a72 100644 --- a/packages/core/src/runtime/session.ts +++ b/packages/core/src/runtime/session.ts @@ -7,18 +7,24 @@ import { withDefaultDeps } from '@memo/core/runtime/defaults' import { parseAssistant } from '@memo/core/utils/utils' import type { CallToolResult } from '@modelcontextprotocol/sdk/types' import type { + ActionPayload, AgentSession, AgentSessionDeps, AgentSessionOptions, AgentStepTrace, ChatMessage, + FinalPayload, HistoryEvent, HistorySink, + Hook, + HookMiddleware, LLMResponse, + ObservationPayload, ParsedAssistant, SessionMode, TokenCounter, TokenUsage, + TurnStartPayload, TurnResult, TurnStatus, } from '@memo/core/types' @@ -61,6 +67,22 @@ async function emitEventToSinks(event: HistoryEvent, sinks: HistorySink[]) { } } +async function runHookChain(hook: Hook | undefined, ...args: TArgs) { + const middlewares = normalizeHook(hook, args.length) + if (!middlewares.length) return + let idx = -1 + const dispatch = async (i: number): Promise => { + if (i <= idx) { + throw new Error('hook middleware 调用顺序错误：next 被重复触发') + } + idx = i + const fn = middlewares[i] + if (!fn) return + await fn(() => dispatch(i + 1), ...args) + } + await dispatch(0) +} + function flattenCallToolResult(result: CallToolResult): string { const texts = result.content?.flatMap((item) => { @@ -94,6 +116,24 @@ function parseToolInput(tool: ToolRegistry[string], rawInput: unknown) { return { ok: true as const, data: parsed.data } } +function normalizeHook( + hook: Hook | undefined, + argsLength: number, +): HookMiddleware[] { + if (!hook) return [] + const list = Array.isArray(hook) ? hook : [hook] + return list.map((fn) => { + // 若函数签名不包含 next，视为普通回调，执行后自动进入下一个 middleware。 + if (fn.length <= argsLength) { + return async (next: () => Promise, ...args: TArgs) => { + await (fn as (...args: TArgs) => unknown)(...args) + await next() + } + } + return fn + }) +} + /** 进程内的对话 Session，实现多轮运行与日志写入。 */ class AgentSessionImpl implements AgentSession { public id: string @@ -154,6 +194,7 @@ class AgentSessionImpl implements AgentSession { content: input, meta: { tokens: { prompt: promptTokens } }, }) + await runHookChain(this.deps.onTurnStart, { turn, input }) let finalText = '' let status: TurnStatus = 'ok' @@ -176,6 +217,12 @@ class AgentSessionImpl implements AgentSession { role: 'assistant', meta: { tokens: { prompt: estimatedPrompt } }, }) + await runHookChain(this.deps.onFinal, { + turn, + step, + final: limitMessage, + status, + }) break } if (this.options.warnPromptTokens && estimatedPrompt > this.options.warnPromptTokens) { @@ -201,6 +248,11 @@ class AgentSessionImpl implements AgentSession { finalText = msg errorMessage = msg await this.emitEvent('final', { turn, content: msg, role: 'assistant' }) + await runHookChain(this.deps.onFinal, { + turn, + final: msg, + status, + }) break } @@ -247,10 +299,22 @@ class AgentSessionImpl implements AgentSession { role: 'assistant', meta: { tokens: stepUsage }, }) + await runHookChain(this.deps.onFinal, { + turn, + step, + final: parsed.final, + status, + }) break } if (parsed.action) { + await runHookChain(this.deps.onAction, { + turn, + step, + tool: parsed.action.tool, + input: parsed.action.input, + }) await this.emitEvent('action', { turn, step, @@ -283,7 +347,12 @@ class AgentSessionImpl implements AgentSession { if (lastStep) { lastStep.observation = observation } - this.deps.onObservation?.(parsed.action.tool, observation, step) + await runHookChain(this.deps.onObservation, { + turn, + step, + tool: parsed.action.tool, + observation, + }) await this.emitEvent('observation', { turn, @@ -311,6 +380,11 @@ class AgentSessionImpl implements AgentSession { content: finalText, role: 'assistant', }) + await runHookChain(this.deps.onFinal, { + turn, + final: finalText, + status, + }) } await this.emitEvent('turn_end', { diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index 7d65c267..4b3d2242 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -95,8 +95,19 @@ export type AgentDeps = { loadPrompt?: () => Promise /** 每次 assistant 输出时的回调。 */ onAssistantStep?: (content: string, step: number) => void - /** 每次工具 observation 返回时的回调。 */ - onObservation?: (tool: string, observation: string, step: number) => void + /** + * Turn 开始的 Hook，可用于接入中间件（日志、链路追踪等）。 + * 支持直接回调或 middleware 链（next -> ...）。 + */ + onTurnStart?: Hook<[TurnStartPayload]> + /** + * 模型决定调用工具时的 Hook，可用于审核工具参数、记录链路等。 + */ + onAction?: Hook<[ActionPayload]> + /** 每次工具 observation 返回时的 Hook，支持中间件。 */ + onObservation?: Hook<[ObservationPayload]> + /** 最终输出时的 Hook，携带状态码，支持中间件。 */ + onFinal?: Hook<[FinalPayload]> /** 资源释放回调（如关闭 MCP Client）。 */ dispose?: () => Promise } @@ -191,6 +202,41 @@ export type HistoryEvent = { meta?: Record } +/** Hook middleware：形如 (next, payload...)，可通过 next 形成链式调用。 */ +export type HookMiddleware = ( + next: () => Promise, + ...args: TArgs +) => void | Promise + +/** 普通 Hook 回调。 */ +export type HookHandler = (...args: TArgs) => void | Promise + +/** + * Hook 类型：支持单个回调、单个 middleware，或数组（多中间件）。 + * - 若函数签名不包含 next，则视为普通回调，自动串联到 middleware 链中。 + */ +export type Hook = + | HookHandler + | HookMiddleware + | Array | HookMiddleware> + +/** Turn 开始的 Hook 入参。 */ +export type TurnStartPayload = { turn: number; input: string } + +/** Action Hook 入参。 */ +export type ActionPayload = { turn: number; step: number; tool: string; input: unknown } + +/** Observation Hook 入参。 */ +export type ObservationPayload = { + turn: number + step: number + tool: string + observation: string +} + +/** Final Hook 入参。 */ +export type FinalPayload = { turn: number; step?: number; final: string; status: TurnStatus } + /** 历史落盘抽象，可输出到文件/外部系统。 */ export type HistorySink = { /** 写入单条事件。 */ diff --git a/packages/ui/src/index.ts b/packages/ui/src/index.ts index 2188ad09..f4ecc409 100644 --- a/packages/ui/src/index.ts +++ b/packages/ui/src/index.ts @@ -116,7 +116,7 @@ async function runInteractive(parsed: ParsedArgs) { } process.stdout.write(text) }, - onObservation: (tool: string, observation: string, step: number) => { + onObservation: ({ tool, observation, step }) => { console.log(`\n[第 ${step + 1} 步 工具=${tool}]\n${observation}\n`) }, }