From a0d031e30a6bc5c5b97198d563dec5c2fc3a3ffc Mon Sep 17 00:00:00 2001 From: MemOS Fix Date: Thu, 12 Mar 2026 19:34:58 +0000 Subject: [PATCH] fix(hooks): expose sessionKey and agentId in agent_end and before_agent_start events This fix allows memory plugins (like memos-local) to properly identify which agent is running when capturing conversation history. Previously, agent_end and before_agent_start hooks only received generic event data without session context. Memory plugins could not distinguish between different agents (e.g., 'main' vs 'qian-duoduo'), causing all memories to be tagged with the wrong owner. Changes: - Add sessionKey and agentId fields to PluginHookAgentEndEvent - Add sessionKey and agentId fields to PluginHookBeforeAgentStartEvent - Pass sessionKey and agentId when triggering agent_end hook - Pass sessionKey and agentId when triggering before_agent_start hook Fixes context issue where multi-agent setups could not properly isolate per-agent memory. --- src/agents/pi-embedded-runner/run.ts | 6 +- src/agents/pi-embedded-runner/run/attempt.ts | 6 +- src/plugins/hooks.agent-context.test.ts | 163 +++++++++++++++++++ src/plugins/types.ts | 8 + 4 files changed, 181 insertions(+), 2 deletions(-) create mode 100644 src/plugins/hooks.agent-context.test.ts diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 0c66203992f..1171797f6d8 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -348,7 +348,11 @@ export async function runEmbeddedPiAgent( if (hookRunner?.hasHooks("before_agent_start")) { try { legacyBeforeAgentStartResult = await hookRunner.runBeforeAgentStart( - { prompt: params.prompt }, + { + prompt: params.prompt, + sessionKey: params.sessionKey, + agentId: workspaceResolution.agentId, + }, hookCtx, ); modelResolveOverride = { diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index d785218f819..080f49efea1 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -157,7 +157,7 @@ type PromptBuildHookRunner = { ctx: PluginHookAgentContext, ) => Promise; runBeforeAgentStart: ( - event: { prompt: string; messages: unknown[] }, + event: { prompt: string; messages?: unknown[]; sessionKey?: string; agentId?: string }, ctx: PluginHookAgentContext, ) => Promise; }; @@ -1439,6 +1439,8 @@ export async function resolvePromptBuildHookResult(params: { { prompt: params.prompt, messages: params.messages, + sessionKey: params.hookCtx.sessionKey, + agentId: params.hookCtx.agentId, }, params.hookCtx, ) @@ -3078,6 +3080,8 @@ export async function runEmbeddedAttempt( success: !aborted && !promptError, error: promptError ? describeUnknownError(promptError) : undefined, durationMs: Date.now() - promptStartedAt, + sessionKey: params.sessionKey, + agentId: hookAgentId, }, { agentId: hookAgentId, diff --git a/src/plugins/hooks.agent-context.test.ts b/src/plugins/hooks.agent-context.test.ts new file mode 100644 index 00000000000..59788be5f2d --- /dev/null +++ b/src/plugins/hooks.agent-context.test.ts @@ -0,0 +1,163 @@ +/** + * Tests for agent context (sessionKey, agentId) in hook events + * + * Validates that sessionKey and agentId are correctly passed to + * agent_end and before_agent_start hook handlers. + */ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { createHookRunner } from "./hooks.js"; +import { addTestHook, TEST_PLUGIN_AGENT_CTX } from "./hooks.test-helpers.js"; +import { createEmptyPluginRegistry, type PluginRegistry } from "./registry.js"; +import type { PluginHookRegistration } from "./types.js"; + +function addAgentEndHook( + registry: PluginRegistry, + pluginId: string, + handler: (event: { messages: unknown[]; success: boolean; sessionKey?: string; agentId?: string }) => void | Promise, + priority?: number, +) { + addTestHook({ + registry, + pluginId, + hookName: "agent_end", + handler: handler as PluginHookRegistration["handler"], + priority, + }); +} + +function addBeforeAgentStartHook( + registry: PluginRegistry, + pluginId: string, + handler: (event: { prompt: string; messages?: unknown[]; sessionKey?: string; agentId?: string }) => void | Promise, + priority?: number, +) { + addTestHook({ + registry, + pluginId, + hookName: "before_agent_start", + handler: handler as PluginHookRegistration["handler"], + priority, + }); +} + +describe("hook events include sessionKey and agentId", () => { + let registry: PluginRegistry; + const stubCtx = TEST_PLUGIN_AGENT_CTX; + + beforeEach(() => { + registry = createEmptyPluginRegistry(); + }); + + describe("agent_end hook", () => { + it("receives sessionKey and agentId in event", async () => { + const handler = vi.fn().mockResolvedValue(undefined); + + addAgentEndHook(registry, "memory-plugin", handler); + + const runner = createHookRunner(registry); + const testMessages = [{ role: "user", content: "hello" }]; + + await runner.runAgentEnd( + { + messages: testMessages, + success: true, + sessionKey: "agent:assistant-beta:main", + agentId: "assistant-beta", + }, + stubCtx, + ); + + expect(handler).toHaveBeenCalledWith( + expect.objectContaining({ + messages: testMessages, + success: true, + sessionKey: "agent:assistant-beta:main", + agentId: "assistant-beta", + }), + stubCtx, + ); + }); + + it("works with default values when sessionKey and agentId are not provided", async () => { + const handler = vi.fn().mockResolvedValue(undefined); + + addAgentEndHook(registry, "memory-plugin", handler); + + const runner = createHookRunner(registry); + const testMessages = [{ role: "user", content: "hello" }]; + + // Call without sessionKey and agentId (backward compatibility) + await runner.runAgentEnd( + { + messages: testMessages, + success: true, + }, + stubCtx, + ); + + expect(handler).toHaveBeenCalledWith( + expect.objectContaining({ + messages: testMessages, + success: true, + sessionKey: undefined, + agentId: undefined, + }), + stubCtx, + ); + }); + }); + + describe("before_agent_start hook", () => { + it("receives sessionKey and agentId in event", async () => { + const handler = vi.fn().mockResolvedValue(undefined); + + addBeforeAgentStartHook(registry, "memory-plugin", handler); + + const runner = createHookRunner(registry); + + await runner.runBeforeAgentStart( + { + prompt: "hello", + messages: [{ role: "user", content: "hello" }], + sessionKey: "agent:assistant-beta:main", + agentId: "assistant-beta", + }, + stubCtx, + ); + + expect(handler).toHaveBeenCalledWith( + expect.objectContaining({ + prompt: "hello", + sessionKey: "agent:assistant-beta:main", + agentId: "assistant-beta", + }), + stubCtx, + ); + }); + + it("works with legacy event shape (backward compatibility)", async () => { + const handler = vi.fn().mockResolvedValue({ prependContext: "context" }); + + addBeforeAgentStartHook(registry, "legacy-plugin", handler); + + const runner = createHookRunner(registry); + + // Call with legacy event shape (without sessionKey and agentId) + await runner.runBeforeAgentStart( + { + prompt: "hello", + }, + stubCtx, + ); + + expect(handler).toHaveBeenCalledWith( + expect.objectContaining({ + prompt: "hello", + sessionKey: undefined, + agentId: undefined, + }), + stubCtx, + ); + }); + }); +}); diff --git a/src/plugins/types.ts b/src/plugins/types.ts index 343a338c4f8..96c54c6d7a0 100644 --- a/src/plugins/types.ts +++ b/src/plugins/types.ts @@ -1513,6 +1513,10 @@ export type PluginHookBeforeAgentStartEvent = { prompt: string; /** Optional because legacy hook can run in pre-session phase. */ messages?: unknown[]; + /** Session key for this agent run (e.g., "agent:qian-duoduo:main") */ + sessionKey?: string; + /** Agent ID for this run (e.g., "qian-duoduo") */ + agentId?: string; }; export type PluginHookBeforeAgentStartResult = PluginHookBeforePromptBuildResult & @@ -1573,6 +1577,10 @@ export type PluginHookAgentEndEvent = { success: boolean; error?: string; durationMs?: number; + /** Session key for this agent run (e.g., "agent:qian-duoduo:main") */ + sessionKey?: string; + /** Agent ID for this run (e.g., "qian-duoduo") */ + agentId?: string; }; // Compaction hooks