feat: add hook middleware support

Author: minorcell

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 均为可选，未提供时不会影响主流程。

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',
+        ])
+    })
+})

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<TArgs extends any[]>(hook: Hook<TArgs> | undefined, ...args: TArgs) {
+    const middlewares = normalizeHook(hook, args.length)
+    if (!middlewares.length) return
+    let idx = -1
+    const dispatch = async (i: number): Promise<void> => {
+        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<TArgs extends any[]>(
+    hook: Hook<TArgs> | undefined,
+    argsLength: number,
+): HookMiddleware<TArgs>[] {
+    if (!hook) return []
+    const list = Array.isArray(hook) ? hook : [hook]
+    return list.map((fn) => {
+        // 若函数签名不包含 next，视为普通回调，执行后自动进入下一个 middleware。
+        if (fn.length <= argsLength) {
+            return async (next: () => Promise<void>, ...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<TurnStartPayload>(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<FinalPayload>(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<FinalPayload>(this.deps.onFinal, {
+                    turn,
+                    final: msg,
+                    status,
+                })
                 break
             }
 
@@ -247,10 +299,22 @@ class AgentSessionImpl implements AgentSession {
                     role: 'assistant',
                     meta: { tokens: stepUsage },
                 })
+                await runHookChain<FinalPayload>(this.deps.onFinal, {
+                    turn,
+                    step,
+                    final: parsed.final,
+                    status,
+                })
                 break
             }
 
             if (parsed.action) {
+                await runHookChain<ActionPayload>(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<ObservationPayload>(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<FinalPayload>(this.deps.onFinal, {
+                turn,
+                final: finalText,
+                status,
+            })
         }
 
         await this.emitEvent('turn_end', {

packages/core/src/types.ts

@@ -95,8 +95,19 @@ export type AgentDeps = {
     loadPrompt?: () => Promise<string>
     /** 每次 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<void>
 }
@@ -191,6 +202,41 @@ export type HistoryEvent = {
     meta?: Record<string, unknown>
 }
 
+/** Hook middleware：形如 (next, payload...)，可通过 next 形成链式调用。 */
+export type HookMiddleware<TArgs extends any[]> = (
+    next: () => Promise<void>,
+    ...args: TArgs
+) => void | Promise<void>
+
+/** 普通 Hook 回调。 */
+export type HookHandler<TArgs extends any[]> = (...args: TArgs) => void | Promise<void>
+
+/**
+ * Hook 类型：支持单个回调、单个 middleware，或数组（多中间件）。
+ * - 若函数签名不包含 next，则视为普通回调，自动串联到 middleware 链中。
+ */
+export type Hook<TArgs extends any[]> =
+    | HookHandler<TArgs>
+    | HookMiddleware<TArgs>
+    | Array<HookHandler<TArgs> | HookMiddleware<TArgs>>
+
+/** 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 = {
     /** 写入单条事件。 */

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`)
         },
     }