Merge 4314f44fe22ec8a99093892d804085c806033bdd into 9fb78453e088cd7b553d7779faa0de5c83708e70

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 = {

src/agents/pi-embedded-runner/run/attempt.test.ts

@@ -114,6 +114,34 @@ describe("resolvePromptBuildHookResult", () => {
     expect(result.prependContext).toBe("from-hook");
   });
 
+  it("includes legacy hook agent context fields only when present", async () => {
+    const hookRunner = createLegacyOnlyHookRunner();
+    const messages = [{ role: "user", content: "ctx" }];
+
+    await resolvePromptBuildHookResult({
+      prompt: "hello",
+      messages,
+      hookCtx: {
+        sessionKey: "agent:assistant-beta:main",
+        agentId: "assistant-beta",
+      },
+      hookRunner,
+    });
+
+    expect(hookRunner.runBeforeAgentStart).toHaveBeenCalledWith(
+      {
+        prompt: "hello",
+        messages,
+        sessionKey: "agent:assistant-beta:main",
+        agentId: "assistant-beta",
+      },
+      {
+        sessionKey: "agent:assistant-beta:main",
+        agentId: "assistant-beta",
+      },
+    );
+  });
+
   it("merges prompt-build and legacy context fields in deterministic order", async () => {
     const hookRunner = {
       hasHooks: vi.fn(() => true),

src/agents/pi-embedded-runner/run/attempt.ts

@@ -157,7 +157,7 @@ type PromptBuildHookRunner = {
     ctx: PluginHookAgentContext,
   ) => Promise<PluginHookBeforePromptBuildResult | undefined>;
   runBeforeAgentStart: (
-    event: { prompt: string; messages: unknown[] },
+    event: { prompt: string; messages?: unknown[]; sessionKey?: string; agentId?: string },
     ctx: PluginHookAgentContext,
   ) => Promise<PluginHookBeforeAgentStartResult | undefined>;
 };
@@ -1439,6 +1439,8 @@ export async function resolvePromptBuildHookResult(params: {
             {
               prompt: params.prompt,
               messages: params.messages,
+              ...(params.hookCtx.sessionKey ? { sessionKey: params.hookCtx.sessionKey } : {}),
+              ...(params.hookCtx.agentId ? { 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,

src/plugins/hooks.agent-context.test.ts

@@ -0,0 +1,167 @@
+/**
+ * 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<void>,
+  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<void>,
+  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).toHaveBeenCalledTimes(1);
+      expect(handler.mock.calls[0]?.[0]).toEqual({
+        messages: testMessages,
+        success: true,
+      });
+      expect(handler.mock.calls[0]?.[1]).toBe(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).toHaveBeenCalledTimes(1);
+      expect(handler.mock.calls[0]?.[0]).toEqual({
+        prompt: "hello",
+      });
+      expect(handler.mock.calls[0]?.[1]).toBe(stubCtx);
+    });
+  });
+});

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