diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 50695c76a..3875421f3 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -11,6 +11,7 @@ - Added `ctx.modelRegistry.stream()` and `streamSimple()` for extension model calls through configured providers with resolved authentication ([#8964](https://github.com/earendil-works/pi/issues/8964)). - Added per-model `reserveTokens` and `keepRecentTokens` settings through `compaction.modelOverrides`, with ordinary compaction settings as fallback ([#8133](https://github.com/earendil-works/pi-mono/issues/8133)). - Added `compat.allowedFallbackModels` configuration for overriding or disabling Anthropic server-side fallback models ([#9294](https://github.com/earendil-works/pi/issues/9294)). +- Added an unsubscribe function from `pi.on()` so extensions can drop event handlers. Handlers added or removed during a dispatch apply to later dispatches, not the current one ([#8967](https://github.com/earendil-works/pi/issues/8967)). ### Changed diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 24241eead..36f4fc2f7 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -1371,7 +1371,16 @@ export default function (pi: ExtensionAPI) { ### pi.on(event, handler) -Subscribe to events. See [Events](#events) for event types and return values. +Subscribe to events. Returns an unsubscribe function that removes only that registration. See [Events](#events) for event types and return values. + +```typescript +const unsubscribe = pi.on("agent_end", async (event) => { + unsubscribe(); + await updateIntegration(event.messages); +}); +``` + +Handlers run in extension load order, then registration order within each extension. Adding or removing a handler does not affect a dispatch already in progress. ### pi.registerTool(definition) diff --git a/packages/coding-agent/src/core/extensions/loader.ts b/packages/coding-agent/src/core/extensions/loader.ts index d06f801a4..56a773ab0 100644 --- a/packages/coding-agent/src/core/extensions/loader.ts +++ b/packages/coding-agent/src/core/extensions/loader.ts @@ -277,11 +277,21 @@ function createExtensionAPI( const api = { // Registration methods - write to extension - on(event: string, handler: HandlerFn): void { + on(event: string, handler: HandlerFn): () => void { assertActive(); + const registeredHandler: HandlerFn = (...args) => handler(...args); const list = extension.handlers.get(event) ?? []; - list.push(handler); + list.push(registeredHandler); extension.handlers.set(event, list); + + return () => { + const handlers = extension.handlers.get(event); + if (!handlers) return; + const handlerIndex = handlers.indexOf(registeredHandler); + if (handlerIndex === -1) return; + handlers.splice(handlerIndex, 1); + if (handlers.length === 0) extension.handlers.delete(event); + }; }, registerTool(tool: ToolDefinition): void { diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 68d008694..91f4aa8fd 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -232,18 +232,19 @@ export async function emitSessionShutdownEvent( return false; } +function snapshotEventHandlers(extensions: Extension[], event: ExtensionEvent["type"]) { + return extensions.map((ext) => ({ ext, handlers: ext.handlers.get(event)?.slice() ?? [] })); +} + export async function emitProjectTrustEvent( extensionsResult: LoadExtensionsResult, event: ProjectTrustEvent, ctx: ProjectTrustContext, ): Promise<{ result?: ProjectTrustEventResult; errors: ExtensionError[] }> { const errors: ExtensionError[] = []; - for (const ext of extensionsResult.extensions) { + for (const { ext, handlers } of snapshotEventHandlers(extensionsResult.extensions, "project_trust")) { // A single extension may register multiple handlers for the same event. // The first project_trust handler that returns yes/no wins; undecided falls through. - const handlers = ext.handlers.get("project_trust"); - if (!handlers || handlers.length === 0) continue; - for (const handler of handlers) { try { const handlerResult = (await handler(event, ctx)) as ProjectTrustEventResult; @@ -885,10 +886,7 @@ export class ExtensionRunner { const ctx = this.createContext(); let result: SessionBeforeEventResult | undefined; - for (const ext of this.extensions) { - const handlers = ext.handlers.get(event.type); - if (!handlers || handlers.length === 0) continue; - + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, event.type)) { for (const handler of handlers) { try { const handlerResult = await handler(event, ctx); @@ -920,10 +918,7 @@ export class ExtensionRunner { let currentMessage = event.message; let modified = false; - for (const ext of this.extensions) { - const handlers = ext.handlers.get("message_end"); - if (!handlers || handlers.length === 0) continue; - + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "message_end")) { for (const handler of handlers) { try { const currentEvent: MessageEndEvent = { ...event, message: currentMessage }; @@ -962,10 +957,7 @@ export class ExtensionRunner { const currentEvent: ToolResultEvent = { ...event }; let modified = false; - for (const ext of this.extensions) { - const handlers = ext.handlers.get("tool_result"); - if (!handlers || handlers.length === 0) continue; - + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "tool_result")) { for (const handler of handlers) { try { const handlerResult = (await handler(currentEvent, ctx)) as ToolResultEventResult | undefined; @@ -1016,10 +1008,7 @@ export class ExtensionRunner { const ctx = this.createContext(); let result: ToolCallEventResult | undefined; - for (const ext of this.extensions) { - const handlers = ext.handlers.get("tool_call"); - if (!handlers || handlers.length === 0) continue; - + for (const { handlers } of snapshotEventHandlers(this.extensions, "tool_call")) { for (const handler of handlers) { const handlerResult = await handler(event, ctx); @@ -1038,10 +1027,7 @@ export class ExtensionRunner { async emitUserBash(event: UserBashEvent): Promise { const ctx = this.createContext(); - for (const ext of this.extensions) { - const handlers = ext.handlers.get("user_bash"); - if (!handlers || handlers.length === 0) continue; - + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "user_bash")) { for (const handler of handlers) { try { const handlerResult = await handler(event, ctx); @@ -1073,10 +1059,7 @@ export class ExtensionRunner { const ctx = this.createContext(); let currentMessages = structuredClone(messages); - for (const ext of this.extensions) { - const handlers = ext.handlers.get("context"); - if (!handlers || handlers.length === 0) continue; - + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "context")) { for (const handler of handlers) { try { const event: ContextEvent = { type: "context", messages: currentMessages }; @@ -1105,10 +1088,7 @@ export class ExtensionRunner { const ctx = this.createContext(); let currentPayload = payload; - for (const ext of this.extensions) { - const handlers = ext.handlers.get("before_provider_request"); - if (!handlers || handlers.length === 0) continue; - + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "before_provider_request")) { for (const handler of handlers) { try { const event: BeforeProviderRequestEvent = { @@ -1138,10 +1118,7 @@ export class ExtensionRunner { async emitBeforeProviderHeaders(headers: ProviderHeaders): Promise { const ctx = this.createContext(); - for (const ext of this.extensions) { - const handlers = ext.handlers.get("before_provider_headers"); - if (!handlers || handlers.length === 0) continue; - + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "before_provider_headers")) { for (const handler of handlers) { try { // Handlers mutate `headers` in place; the return value is ignored. @@ -1183,10 +1160,7 @@ export class ExtensionRunner { }; const messages: NonNullable[] = []; - for (const ext of this.extensions) { - const handlers = ext.handlers.get("before_agent_start"); - if (!handlers || handlers.length === 0) continue; - + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "before_agent_start")) { for (const handler of handlers) { try { const event: BeforeAgentStartEvent = { @@ -1236,10 +1210,7 @@ export class ExtensionRunner { const promptPaths: Array<{ path: string; extensionPath: string }> = []; const themePaths: Array<{ path: string; extensionPath: string }> = []; - for (const ext of this.extensions) { - const handlers = ext.handlers.get("resources_discover"); - if (!handlers || handlers.length === 0) continue; - + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "resources_discover")) { for (const handler of handlers) { try { const event: ResourcesDiscoverEvent = { type: "resources_discover", cwd, reason }; @@ -1282,8 +1253,8 @@ export class ExtensionRunner { let currentText = text; let currentImages = images; - for (const ext of this.extensions) { - for (const handler of ext.handlers.get("input") ?? []) { + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, "input")) { + for (const handler of handlers) { try { const event: InputEvent = { type: "input", diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index afbd82b74..986ae6152 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -1259,51 +1259,63 @@ export interface ExtensionAPI { // Event Subscription // ========================================================================= - on(event: "project_trust", handler: ProjectTrustHandler): void; - on(event: "resources_discover", handler: ExtensionHandler): void; - on(event: "session_start", handler: ExtensionHandler): void; - on(event: "session_info_changed", handler: ExtensionHandler): void; + on(event: "project_trust", handler: ProjectTrustHandler): () => void; + on( + event: "resources_discover", + handler: ExtensionHandler, + ): () => void; + on(event: "session_start", handler: ExtensionHandler): () => void; + on(event: "session_info_changed", handler: ExtensionHandler): () => void; on( event: "session_before_switch", handler: ExtensionHandler, - ): void; - on(event: "session_before_fork", handler: ExtensionHandler): void; + ): () => void; + on( + event: "session_before_fork", + handler: ExtensionHandler, + ): () => void; on( event: "session_before_compact", handler: ExtensionHandler, - ): void; - on(event: "session_compact", handler: ExtensionHandler): void; - on(event: "session_compact_failed", handler: ExtensionHandler): void; - on(event: "session_shutdown", handler: ExtensionHandler): void; - on(event: "session_before_tree", handler: ExtensionHandler): void; - on(event: "session_tree", handler: ExtensionHandler): void; - on(event: "context", handler: ExtensionHandler): void; + ): () => void; + on(event: "session_compact", handler: ExtensionHandler): () => void; + on(event: "session_compact_failed", handler: ExtensionHandler): () => void; + on(event: "session_shutdown", handler: ExtensionHandler): () => void; + on( + event: "session_before_tree", + handler: ExtensionHandler, + ): () => void; + on(event: "session_tree", handler: ExtensionHandler): () => void; + on(event: "context", handler: ExtensionHandler): () => void; on( event: "before_provider_request", handler: ExtensionHandler, - ): void; - on(event: "before_provider_headers", handler: ExtensionHandler): void; - on(event: "after_provider_response", handler: ExtensionHandler): void; - on(event: "before_agent_start", handler: ExtensionHandler): void; - on(event: "agent_start", handler: ExtensionHandler): void; - on(event: "agent_end", handler: ExtensionHandler): void; - on(event: "agent_settled", handler: ExtensionHandler): void; - on(event: "ui_prompt_start", handler: ExtensionHandler): void; - on(event: "ui_prompt_end", handler: ExtensionHandler): void; - on(event: "turn_start", handler: ExtensionHandler): void; - on(event: "turn_end", handler: ExtensionHandler): void; - on(event: "message_start", handler: ExtensionHandler): void; - on(event: "message_update", handler: ExtensionHandler): void; - on(event: "message_end", handler: ExtensionHandler): void; - on(event: "tool_execution_start", handler: ExtensionHandler): void; - on(event: "tool_execution_update", handler: ExtensionHandler): void; - on(event: "tool_execution_end", handler: ExtensionHandler): void; - on(event: "model_select", handler: ExtensionHandler): void; - on(event: "thinking_level_select", handler: ExtensionHandler): void; - on(event: "tool_call", handler: ExtensionHandler): void; - on(event: "tool_result", handler: ExtensionHandler): void; - on(event: "user_bash", handler: ExtensionHandler): void; - on(event: "input", handler: ExtensionHandler): void; + ): () => void; + on(event: "before_provider_headers", handler: ExtensionHandler): () => void; + on(event: "after_provider_response", handler: ExtensionHandler): () => void; + on( + event: "before_agent_start", + handler: ExtensionHandler, + ): () => void; + on(event: "agent_start", handler: ExtensionHandler): () => void; + on(event: "agent_end", handler: ExtensionHandler): () => void; + on(event: "agent_settled", handler: ExtensionHandler): () => void; + on(event: "ui_prompt_start", handler: ExtensionHandler): () => void; + on(event: "ui_prompt_end", handler: ExtensionHandler): () => void; + on(event: "turn_start", handler: ExtensionHandler): () => void; + on(event: "turn_end", handler: ExtensionHandler): () => void; + on(event: "message_start", handler: ExtensionHandler): () => void; + on(event: "message_update", handler: ExtensionHandler): () => void; + on(event: "message_end", handler: ExtensionHandler): () => void; + on(event: "tool_execution_start", handler: ExtensionHandler): () => void; + on(event: "tool_execution_update", handler: ExtensionHandler): () => void; + on(event: "tool_execution_end", handler: ExtensionHandler): () => void; + on(event: "model_select", handler: ExtensionHandler): () => void; + on(event: "thinking_level_select", handler: ExtensionHandler): () => void; + on(event: "tool_call", handler: ExtensionHandler): () => void; + on(event: "tool_result", handler: ExtensionHandler): () => void; + on(event: "user_bash", handler: ExtensionHandler): () => void; + on(event: "input", handler: ExtensionHandler): () => void; // ========================================================================= // Tool Registration diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 88088685d..e77b5fd08 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -8,11 +8,18 @@ import * as os from "node:os"; import * as path from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import { createExtensionRuntime, discoverAndLoadExtensions, loadExtensions } from "../src/core/extensions/loader.ts"; +import { createEventBus } from "../src/core/event-bus.ts"; +import { + createExtensionRuntime, + discoverAndLoadExtensions, + loadExtensionFromFactory, + loadExtensions, +} from "../src/core/extensions/loader.ts"; import { ExtensionRunner, emitProjectTrustEvent } from "../src/core/extensions/runner.ts"; import type { ExtensionActions, ExtensionContextActions, + ExtensionFactory, ExtensionUIContext, ProviderConfig, } from "../src/core/extensions/types.ts"; @@ -1068,6 +1075,127 @@ describe("ExtensionRunner", () => { }); }); + // #8967: event handler unsubscription must not disturb other registrations. + describe("event subscriptions", () => { + async function loadSubscriptionExtension(factory: ExtensionFactory) { + const runtime = createExtensionRuntime(); + const extension = await loadExtensionFromFactory(factory, tempDir, createEventBus(), runtime); + const runner = new ExtensionRunner([extension], runtime, tempDir, sessionManager, modelRegistry); + return { extension, runner }; + } + + it("allows self-removal without skipping neighboring handlers", async () => { + const calls: string[] = []; + const { runner } = await loadSubscriptionExtension((pi) => { + const unsubscribe = pi.on("agent_end", () => { + calls.push("A"); + unsubscribe(); + }); + pi.on("agent_end", () => { + calls.push("B"); + }); + }); + + await runner.emit({ type: "agent_end", messages: [] }); + expect(calls).toEqual(["A", "B"]); + + await runner.emit({ type: "agent_end", messages: [] }); + expect(calls).toEqual(["A", "B", "B"]); + }); + + it("removes duplicate registrations independently and cleans up the last handler", async () => { + const calls: string[] = []; + const unsubscribers: Array<() => void> = []; + const { extension, runner } = await loadSubscriptionExtension((pi) => { + const shared = () => { + calls.push("shared"); + }; + unsubscribers.push(pi.on("agent_end", shared)); + unsubscribers.push( + pi.on("agent_end", () => { + calls.push("B"); + }), + ); + unsubscribers.push(pi.on("agent_end", shared)); + }); + const [stopFirst, stopB, stopSecond] = unsubscribers; + + stopSecond(); + stopSecond(); + await runner.emit({ type: "agent_end", messages: [] }); + expect(calls).toEqual(["shared", "B"]); + + stopFirst(); + await runner.emit({ type: "agent_end", messages: [] }); + expect(calls).toEqual(["shared", "B", "B"]); + + stopB(); + expect(extension.handlers.has("agent_end")).toBe(false); + }); + + it("keeps removed pending handlers in the current dispatch", async () => { + const calls: string[] = []; + const { runner } = await loadSubscriptionExtension((pi) => { + pi.on("agent_end", () => { + calls.push("A"); + stopB(); + }); + const stopB = pi.on("agent_end", () => { + calls.push("B"); + }); + pi.on("agent_end", () => { + calls.push("C"); + }); + }); + + await runner.emit({ type: "agent_end", messages: [] }); + expect(calls).toEqual(["A", "B", "C"]); + await runner.emit({ type: "agent_end", messages: [] }); + expect(calls).toEqual(["A", "B", "C", "A", "C"]); + }); + + it("defers registrations made during dispatch until the next dispatch", async () => { + const calls: string[] = []; + const { runner } = await loadSubscriptionExtension((pi) => { + pi.on("agent_end", () => { + calls.push("A"); + pi.on("agent_end", () => { + calls.push("C"); + }); + }); + pi.on("agent_end", () => { + calls.push("B"); + }); + }); + + await runner.emit({ type: "agent_end", messages: [] }); + expect(calls).toEqual(["A", "B"]); + await runner.emit({ type: "agent_end", messages: [] }); + expect(calls).toEqual(["A", "B", "A", "B", "C"]); + }); + + it("uses a fresh handler list for nested dispatches", async () => { + const calls: string[] = []; + const { runner } = await loadSubscriptionExtension((pi) => { + const stopA = pi.on("agent_end", async () => { + calls.push("A"); + stopA(); + stopB(); + pi.on("agent_end", () => { + calls.push("C"); + }); + await runner.emit({ type: "agent_end", messages: [] }); + }); + const stopB = pi.on("agent_end", () => { + calls.push("B"); + }); + }); + + await runner.emit({ type: "agent_end", messages: [] }); + expect(calls).toEqual(["A", "C", "B"]); + }); + }); + describe("hasHandlers", () => { it("returns true when handlers exist for event type", async () => { const extCode = `