@zhushanwen/pi-msg-id-mapper 1.0.1

package/index.ts

@@ -0,0 +1 @@
+export { default } from "./src/index.ts";

package/package.json

@@ -0,0 +1,35 @@
+{
+  "name": "@zhushanwen/pi-msg-id-mapper",
+  "version": "1.0.1",
+  "description": "Client UUID ↔ user entry ID mapping extension for Pi — enables metadata preservation across sessions",
+  "type": "module",
+  "main": "src/index.ts",
+  "pi": {
+    "extensions": [
+      "./index.ts"
+    ]
+  },
+  "keywords": [
+    "pi-package",
+    "extension",
+    "msg-id",
+    "client-uuid",
+    "mapping"
+  ],
+  "license": "MIT",
+  "files": [
+    "src/",
+    "index.ts"
+  ],
+  "peerDependencies": {
+    "@earendil-works/pi-coding-agent": "*"
+  },
+  "devDependencies": {
+    "@types/node": "^24.0.0",
+    "vitest": "^4.1.8"
+  },
+  "scripts": {
+    "typecheck": "npx tsc --noEmit",
+    "test": "vitest run"
+  }
+}

package/src/__tests__/msg-id-mapper.test.ts

@@ -0,0 +1,281 @@
+/**
+ * msg-id-mapper extension 真实行为测试。
+ *
+ * 覆盖（替换原 expect(true) 占位）：
+ * - input hook：clientUuid 提取（有/无/畸形 userEntryId 标记、非 rpc source、多标记全剥离）
+ * - message_end：仅 user role + pendingClientUuid 时置 awaiting flag
+ * - flush（message_start/turn_end/agent_end 三重安全网）：映射写入、幂等清空、
+ *   leafId 未就绪重试、异常兜底（不抛出 + 保留 flag 供下次幂等重试）
+ *
+ * mock 策略：pi SDK 经 workspace node_modules 提供类型（import type 零运行时解析），
+ * ExtensionAPI/ExtensionContext 用结构化桩（参照 subagent-workflow mocks/ 的最小桩模式，
+ * 本包无 SDK 运行时依赖故不需要 alias 配置）。
+ *
+ * 运行：cd extensions/msg-id-mapper && npx vitest run
+ */
+import { describe, it, expect, vi, afterEach } from 'vitest'
+import createMapper from '../index'
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+  InputEvent,
+  MessageEndEvent,
+  MessageStartEvent,
+  TurnEndEvent,
+  AgentEndEvent,
+} from '@earendil-works/pi-coding-agent'
+
+const ENTRY_TYPE = 'xyz.client-msg-id'
+
+/** 合法 uuid：`u-` + 36 字符（8-4-4-4-12 hex+连字符），与 TAG_MATCH 严格匹配 */
+const UUID = 'u-123e4567-e89b-12d3-a456-426614174000'
+const marker = (uuid: string): string => `<!--xyz:msg:${uuid}-->`
+
+/** hook 注册表 + 可控桩（每个测试用例通过工厂重建，隔离闭包状态） */
+interface Harness {
+  appendEntry: ReturnType<typeof vi.fn>
+  input: (text: string, source?: string) => unknown
+  /** 传 event 对象本身（供畸形 event 兜底测试直接给非 InputEvent 输入） */
+  inputRaw: (event: unknown) => unknown
+  messageEnd: (role: string) => unknown
+  /** 触发 flush 类 hook（message_start/turn_end/agent_end），ctx.getLeafId 返回 leafId */
+  flush: (event: 'message_start' | 'turn_end' | 'agent_end', leafId?: string | null) => void
+  /** ctx.getLeafId 为可注入实现（供 throw 场景） */
+  flushWithLeafIdFn: (event: 'message_start' | 'turn_end' | 'agent_end', getLeafId: () => string | null) => void
+}
+
+function createHarness(): Harness {
+  const handlers = new Map<string, (...args: unknown[]) => unknown>()
+  const appendEntry = vi.fn()
+  const pi = {
+    on: (event: string, handler: (...args: unknown[]) => unknown) => {
+      handlers.set(event, handler)
+    },
+    appendEntry,
+  } as unknown as ExtensionAPI
+  createMapper(pi)
+
+  const makeCtx = (getLeafId: () => string | null): ExtensionContext =>
+    ({ sessionManager: { getLeafId } }) as unknown as ExtensionContext
+
+  return {
+    appendEntry,
+    input: (text: string, source = 'rpc') =>
+      handlers.get('input')!({ type: 'input', text, source } satisfies InputEvent as unknown as InputEvent),
+    inputRaw: (event: unknown) => handlers.get('input')!(event),
+    messageEnd: (role: string) =>
+      handlers.get('message_end')!({ type: 'message_end', message: { role } } as MessageEndEvent),
+    flush: (event, leafId = 'e-user-entry-1') => {
+      const ctx = makeCtx(() => leafId)
+      if (event === 'message_start') handlers.get(event)!({ type: event } as MessageStartEvent, ctx)
+      if (event === 'turn_end') handlers.get(event)!({ type: event } as TurnEndEvent, ctx)
+      if (event === 'agent_end') handlers.get(event)!({ type: event } as AgentEndEvent, ctx)
+    },
+    flushWithLeafIdFn: (event, getLeafId) => {
+      const ctx = makeCtx(getLeafId)
+      if (event === 'message_start') handlers.get(event)!({ type: event } as MessageStartEvent, ctx)
+      if (event === 'turn_end') handlers.get(event)!({ type: event } as TurnEndEvent, ctx)
+      if (event === 'agent_end') handlers.get(event)!({ type: event } as AgentEndEvent, ctx)
+    },
+  }
+}
+
+afterEach(() => {
+  vi.restoreAllMocks()
+})
+
+describe('input hook · clientUuid 提取（extractClientUuid）', () => {
+  it('rpc + 合法标记 → transform 剥离标记，uuid 进入 pending（后续 flush 可写入映射）', () => {
+    const h = createHarness()
+    const result = h.input(`do the task ${marker(UUID)}`)
+
+    expect(result).toEqual({ action: 'transform', text: 'do the task' })
+
+    // pending 已捕获：走完 message_end + flush 应写出该 uuid 的映射
+    h.messageEnd('user')
+    h.flush('message_start')
+    expect(h.appendEntry).toHaveBeenCalledWith(ENTRY_TYPE, {
+      clientUuid: UUID,
+      userEntryId: 'e-user-entry-1',
+    })
+  })
+
+  it('rpc + 无标记 → continue（文本不动，无映射可写）', () => {
+    const h = createHarness()
+    expect(h.input('plain prompt without marker')).toEqual({ action: 'continue' })
+
+    h.messageEnd('user')
+    h.flush('message_start')
+    expect(h.appendEntry).not.toHaveBeenCalled()
+  })
+
+  it('非 rpc source（interactive）+ 标记 → continue（不提取不剥离）', () => {
+    const h = createHarness()
+    expect(h.input(`interactive prompt ${marker(UUID)}`, 'interactive')).toEqual({ action: 'continue' })
+
+    h.messageEnd('user')
+    h.flush('message_start')
+    expect(h.appendEntry).not.toHaveBeenCalled()
+  })
+
+  it('畸形 uuid（长度不足 36 / 非 hex 字符）→ 标记不匹配，continue 原文透传', () => {
+    const h = createHarness()
+    // 长度不足：u- + 短 hex
+    expect(h.input(`short ${marker('u-abc')}`)).toEqual({ action: 'continue' })
+    // 非 hex 字符（xyz 不在 [0-9a-fA-F-] 内）
+    expect(h.input(`badchar ${marker('u-123e4567e89b12d3a456426614174xyz')}`)).toEqual({
+      action: 'continue',
+    })
+
+    h.messageEnd('user')
+    h.flush('message_start')
+    expect(h.appendEntry).not.toHaveBeenCalled()
+  })
+
+  it('多标记残留 → 全部剥离（TAG_STRIP 全局替换，非只剥第一个）', () => {
+    const h = createHarness()
+    const result = h.input(`task ${marker(UUID)}${marker('u-123e4567-e89b-12d3-a456-426614174111')}`)
+    // 首个 uuid 进 pending
+    expect(result).toEqual({ action: 'transform', text: 'task' })
+    h.messageEnd('user')
+    h.flush('message_start')
+    expect(h.appendEntry).toHaveBeenCalledWith(ENTRY_TYPE, {
+      clientUuid: UUID,
+      userEntryId: 'e-user-entry-1',
+    })
+  })
+
+  it('畸形 event（handler 内抛错）→ 兜底 return undefined 不外抛（console.error 可观测）', () => {
+    const h = createHarness()
+    const errSpy = vi.spyOn(console, 'error').mockImplementation(() => {})
+    // event 为 null → 取 event.source 即 throw，走 catch 兜底
+    expect(h.inputRaw(null)).toBeUndefined()
+    expect(errSpy).toHaveBeenCalledWith(
+      '[xyz-client-msg-id-mapper] input hook error:',
+      expect.any(TypeError),
+    )
+  })
+})
+
+describe('message_end hook · awaiting flag 门控', () => {
+  it('assistant message_end → 不置 flag，flush 不写映射', () => {
+    const h = createHarness()
+    h.input(`task ${marker(UUID)}`)
+    h.messageEnd('assistant')
+    h.flush('message_start')
+    expect(h.appendEntry).not.toHaveBeenCalled()
+  })
+
+  it('user message_end 但无 pendingClientUuid（无前序标记输入）→ flush 不写映射', () => {
+    const h = createHarness()
+    h.messageEnd('user')
+    h.flush('message_start')
+    expect(h.appendEntry).not.toHaveBeenCalled()
+  })
+
+  it('user message_end 但 pendingClientUuid 已被上一轮 flush 清空 → 不置 flag', () => {
+    const h = createHarness()
+    // 第一轮完整走完（写入并清空 pending）
+    h.input(`task ${marker(UUID)}`)
+    h.messageEnd('user')
+    h.flush('message_start')
+    expect(h.appendEntry).toHaveBeenCalledTimes(1)
+
+    // 第二轮：无新标记输入，仅 user message_end → flush 不重复写
+    h.messageEnd('user')
+    h.flush('message_start')
+    expect(h.appendEntry).toHaveBeenCalledTimes(1)
+  })
+})
+
+describe('flush · 映射写入 / 幂等 / 重试 / 异常兜底（writeMapping）', () => {
+  it('主路径：message_start 拿 leafId 写映射（clientUuid ↔ userEntryId）', () => {
+    const h = createHarness()
+    h.input(`task ${marker(UUID)}`)
+    h.messageEnd('user')
+    h.flush('message_start', 'e-leaf-user-42')
+    expect(h.appendEntry).toHaveBeenCalledTimes(1)
+    expect(h.appendEntry).toHaveBeenCalledWith(ENTRY_TYPE, {
+      clientUuid: UUID,
+      userEntryId: 'e-leaf-user-42',
+    })
+  })
+
+  it('幂等清空：flush 后 pending 清空，重复触发（turn_end/agent_end 兜底再打）不再写', () => {
+    const h = createHarness()
+    h.input(`task ${marker(UUID)}`)
+    h.messageEnd('user')
+    h.flush('message_start')
+    expect(h.appendEntry).toHaveBeenCalledTimes(1)
+
+    h.flush('turn_end')
+    h.flush('agent_end')
+    expect(h.appendEntry).toHaveBeenCalledTimes(1)
+  })
+
+  it('abort 兜底：message_start 不来，turn_end 单独完成写入', () => {
+    const h = createHarness()
+    h.input(`task ${marker(UUID)}`)
+    h.messageEnd('user')
+    h.flush('turn_end')
+    expect(h.appendEntry).toHaveBeenCalledTimes(1)
+  })
+
+  it('leafId 未就绪（null）→ 本次不写不抛错，pending 保留，下一 hook 重试成功', () => {
+    const h = createHarness()
+    h.input(`task ${marker(UUID)}`)
+    h.messageEnd('user')
+    // message_start 时 leafId 尚未更新（null）
+    h.flush('message_start', null)
+    expect(h.appendEntry).not.toHaveBeenCalled()
+
+    // turn_end 时 leafId 就绪 → 补写
+    h.flush('turn_end', 'e-leaf-late-1')
+    expect(h.appendEntry).toHaveBeenCalledTimes(1)
+    expect(h.appendEntry).toHaveBeenCalledWith(ENTRY_TYPE, {
+      clientUuid: UUID,
+      userEntryId: 'e-leaf-late-1',
+    })
+  })
+
+  it('appendEntry 抛错 → flush 吞错不外抛，flag 未清，下次 hook 幂等重试', () => {
+    const h = createHarness()
+    const errSpy = vi.spyOn(console, 'error').mockImplementation(() => {})
+    h.appendEntry.mockImplementation(() => {
+      throw new Error('jsonl write failed')
+    })
+
+    h.input(`task ${marker(UUID)}`)
+    h.messageEnd('user')
+    expect(() => h.flush('message_start')).not.toThrow()
+    expect(errSpy).toHaveBeenCalledWith('[xyz-client-msg-id-mapper] flush error:', expect.any(Error))
+
+    // 重试成功：appendEntry 恢复 → turn_end 补写（flag 未清的幂等重试语义）
+    h.appendEntry.mockImplementation(() => {})
+    h.flush('turn_end')
+    expect(h.appendEntry).toHaveBeenCalledTimes(2) // 1 次失败 + 1 次成功
+  })
+
+  it('getLeafId 抛错 → flush 吞错不外抛（catch 兜底覆盖 writeMapping 全程）', () => {
+    const h = createHarness()
+    const errSpy = vi.spyOn(console, 'error').mockImplementation(() => {})
+    h.input(`task ${marker(UUID)}`)
+    h.messageEnd('user')
+
+    expect(() =>
+      h.flushWithLeafIdFn('message_start', () => {
+        throw new Error('session manager gone')
+      }),
+    ).not.toThrow()
+    expect(errSpy).toHaveBeenCalledWith('[xyz-client-msg-id-mapper] flush error:', expect.any(Error))
+  })
+
+  it('无 user message_end（未置 flag）→ 三重安全网全不写', () => {
+    const h = createHarness()
+    h.input(`task ${marker(UUID)}`) // 只有 pending，没有 user 落盘
+    h.flush('message_start')
+    h.flush('turn_end')
+    h.flush('agent_end')
+    expect(h.appendEntry).not.toHaveBeenCalled()
+  })
+})

package/src/index.ts

@@ -0,0 +1,122 @@
+/**
+ * Client UUID ↔ user entry ID mapping extension for Pi.
+ *
+ * Establishes clientUuid ↔ userEntryId mapping for xyz-agent session metadata
+ * preservation across sessions.
+ *
+ * Mechanism (no pi source modification, pure extension hooks):
+ *  1. xyz-agent appends HTML comment marker `<!--xyz:msg:<uuid>-->` to prompt text
+ *     (uuid is the full user message id from appendUser, formatted as `u-<36hex>`).
+ *  2. `input` hook intercepts prompts with source==='rpc', strips the marker
+ *     (LLM doesn't see the transformed content beyond the transform), and stores
+ *     the uuid in pendingClientUuid.
+ *  3. After user `message_end`, the next triggered hook (message_start/turn_end/agent_end)
+ *     reads `ctx.sessionManager.getLeafId()` which now returns the persisted userEntryId.
+ *  4. `pi.appendEntry("xyz.client-msg-id", {clientUuid, userEntryId})` writes the mapping
+ *     into pi JSONL (CustomEntry, not in LLM context).
+ *
+ * Mapping is automatically preserved with fork/clone (in the same JSONL).
+ * When xyz-agent reopens a session, it scans customType==="xyz.client-msg-id" entries
+ * to rebuild the mapping table.
+ *
+ * Degradation strategy: any hook error → pi runner's try/catch swallows it,
+ * mapping is missing, and xyz-agent degrades to textToSegments (split by plain text).
+ * Missing mapping doesn't affect the agent main flow.
+ *
+ * Same package shape as @zhushanwen/pi-system-prompt (extensions/system-prompt/):
+ * TypeScript source with no build step and no runtime deps (peer dep on pi only),
+ * bundled as a builtin extension by scripts/bundle-extensions.mjs.
+ */
+
+import type { ExtensionAPI, ExtensionContext, InputEvent, MessageEndEvent, MessageStartEvent, TurnEndEvent, AgentEndEvent } from '@earendil-works/pi-coding-agent'
+
+// uuid = appendUser 生成的 user message id（`u-` + 36 字符 hex uuid，共 38 字符）。
+// 与 shared SegmentsMetadataEntry.clientUuid 严格一致（同 appendUser 返回值）：
+// extension 写入 custom entry 的 clientUuid = segments.json 的 clientUuid key，
+// entry-tree-builder 据此精确回填 segments（无前缀转换、无双 source of truth）。
+//
+// m1 修复：拆两个正则。TAG_MATCH 非全局，用于 .match 提取第一个 uuid 捕获组（input hook
+// 只需首个 pending uuid）；TAG_STRIP 全局，用于 .replace 剥离所有残留标记（防并发/重试
+// 场景多个标记拼接在同一 prompt 末尾时只剥掉第一个）。两正则分离避免 g 标志的 lastIndex
+// 状态污染（String.match 与全局 exec 混用易踩坑）。
+const TAG_MATCH = /<!--xyz:msg:(u-[0-9a-fA-F-]{36})-->/
+const TAG_STRIP = /<!--xyz:msg:u-[0-9a-fA-F-]{36}-->/g
+const ENTRY_TYPE = 'xyz.client-msg-id'
+
+/** 提取 prompt 里首个标记的 client uuid；无标记 → undefined。 */
+function extractClientUuid(text: string): string | undefined {
+  return text.match(TAG_MATCH)?.[1]
+}
+
+export default function (pi: ExtensionAPI): void {
+  // 待处理的 client uuid（input hook 抓到标记后写入，flush 后清空）。
+  let pendingClientUuid: string | undefined = undefined
+  // user message 已 message_end、等待 flush（拿到 leafId 后写映射）。
+  let awaitingUserPersist = false
+
+  // input hook：拦截 xyz-agent 发来的 prompt（source === 'rpc'），剥离标记。
+  // 非 rpc（interactive/extension）输入不含标记，直接 continue。
+  pi.on('input', (event: InputEvent) => {
+    try {
+      if (event.source !== 'rpc') return { action: 'continue' as const }
+      const clientUuid = extractClientUuid(event.text)
+      if (!clientUuid) return { action: 'continue' as const }
+      pendingClientUuid = clientUuid
+      // transform 后 LLM 看到的是剥离了标记的纯文本。
+      // TAG_STRIP 全局替换：防多个标记残留时只剥掉第一个（m1 修复）。
+      return { action: 'transform' as const, text: event.text.replace(TAG_STRIP, '').trimEnd() }
+    } catch (err) {
+      // 吞错，不阻断主流程（pi runner 也会 try/catch，但显式兜底避免意外 return）。
+      console.error('[xyz-client-msg-id-mapper] input hook error:', err)
+      return undefined
+    }
+  })
+
+  // message_end：user message 持久化前触发。此时 getLeafId() 还指向上一条 entry
+  // （user message 尚未落盘），所以只置 flag，真正 flush 在下一个 hook。
+  pi.on('message_end', (event: MessageEndEvent) => {
+    try {
+      if (event.message.role === 'user' && pendingClientUuid) {
+        awaitingUserPersist = true
+      }
+    } catch (err) {
+      // best-effort：置 flag 失败不阻断消息流——仅丢失该条 clientUuid↔entryId 映射，对话不受影响。
+      console.error('[xyz-client-msg-id-mapper] message_end hook error:', err)
+    }
+    return undefined
+  })
+
+  /** 写映射并清空 pending 状态。leafId 未就绪时静默返回，等下一个 hook 重试。 */
+  const writeMapping = (ctx: ExtensionContext): void => {
+    if (!pendingClientUuid) return
+    // 类型已保证 sessionManager.getLeafId 必有（pi 0.84.1 ReadonlySessionManager，
+    // 返回 string | null）；运行时异常由 flush 的 catch 兜底。
+    const userEntryId = ctx.sessionManager.getLeafId()
+    if (!userEntryId) return // leafId 还没更新，等下一个 hook
+    pi.appendEntry(ENTRY_TYPE, {
+      clientUuid: pendingClientUuid,
+      userEntryId,
+    })
+    pendingClientUuid = undefined
+    awaitingUserPersist = false
+  }
+
+  // flush：读 getLeafId()（= userEntryId，user message 已持久化）+ appendEntry 写映射。
+  // 幂等：写完即清空 pendingClientUuid / awaitingUserPersist，重复触发无副作用。
+  const flush = (ctx: ExtensionContext): void => {
+    if (!awaitingUserPersist) return
+    try {
+      writeMapping(ctx)
+    } catch (err) {
+      // best-effort：映射写入失败不阻断消息流——丢的是本条映射，下次 hook 会因 flag 未清而幂等重试。
+      console.error('[xyz-client-msg-id-mapper] flush error:', err)
+    }
+  }
+
+  // 三重安全网：user message_end 后第一个触发的 hook 拿到的 leafId 才是 userEntryId。
+  // message_start（assistant 开始，此时 user message 已落盘）—— 主路径。
+  pi.on('message_start', (_event: MessageStartEvent, ctx: ExtensionContext) => flush(ctx))
+  // turn_end / agent_end 兜底（abort 等场景 message_start 可能不来）。
+  pi.on('turn_end', (_event: TurnEndEvent, ctx: ExtensionContext) => flush(ctx))
+  pi.on('agent_end', (_event: AgentEndEvent, ctx: ExtensionContext) => flush(ctx))
+}