sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

feat(coding-agent): add event handler unsubscribe (#9630)

Browse source
This commit is contained in:
Ramiz Wachtler 2026-09-17 21:44:26 +02:00 • committed by GitHub
parent 16292398af
commit 46c9de402b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 217 additions and 86 deletions
Download patch file Download diff file Expand all files Collapse all files

1
packages/coding-agent/CHANGELOG.md
View file

@ -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

11
packages/coding-agent/docs/extensions.md
View file

@ -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)

14
packages/coding-agent/src/core/extensions/loader.ts
View file

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

63
packages/coding-agent/src/core/extensions/runner.ts
View file

@ -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<UserBashEventResult | undefined> {
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<ProviderHeaders> {
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<BeforeAgentStartEventResult["message"]>[] = [];
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",

84
packages/coding-agent/src/core/extensions/types.ts
View file

@ -1259,51 +1259,63 @@ export interface ExtensionAPI {
// Event Subscription
// =========================================================================
on(event: "project_trust", handler: ProjectTrustHandler): void;
on(event: "resources_discover", handler: ExtensionHandler<ResourcesDiscoverEvent, ResourcesDiscoverResult>): void;
on(event: "session_start", handler: ExtensionHandler<SessionStartEvent>): void;
on(event: "session_info_changed", handler: ExtensionHandler<SessionInfoChangedEvent>): void;
on(event: "project_trust", handler: ProjectTrustHandler): () => void;
on(
event: "resources_discover",
handler: ExtensionHandler<ResourcesDiscoverEvent, ResourcesDiscoverResult>,
): () => void;
on(event: "session_start", handler: ExtensionHandler<SessionStartEvent>): () => void;
on(event: "session_info_changed", handler: ExtensionHandler<SessionInfoChangedEvent>): () => void;
on(
event: "session_before_switch",
handler: ExtensionHandler<SessionBeforeSwitchEvent, SessionBeforeSwitchResult>,
): void;
on(event: "session_before_fork", handler: ExtensionHandler<SessionBeforeForkEvent, SessionBeforeForkResult>): void;
): () => void;
on(
event: "session_before_fork",
handler: ExtensionHandler<SessionBeforeForkEvent, SessionBeforeForkResult>,
): () => void;
on(
event: "session_before_compact",
handler: ExtensionHandler<SessionBeforeCompactEvent, SessionBeforeCompactResult>,
): void;
on(event: "session_compact", handler: ExtensionHandler<SessionCompactEvent>): void;
on(event: "session_compact_failed", handler: ExtensionHandler<SessionCompactFailedEvent>): void;
on(event: "session_shutdown", handler: ExtensionHandler<SessionShutdownEvent>): void;
on(event: "session_before_tree", handler: ExtensionHandler<SessionBeforeTreeEvent, SessionBeforeTreeResult>): void;
on(event: "session_tree", handler: ExtensionHandler<SessionTreeEvent>): void;
on(event: "context", handler: ExtensionHandler<ContextEvent, ContextEventResult>): void;
): () => void;
on(event: "session_compact", handler: ExtensionHandler<SessionCompactEvent>): () => void;
on(event: "session_compact_failed", handler: ExtensionHandler<SessionCompactFailedEvent>): () => void;
on(event: "session_shutdown", handler: ExtensionHandler<SessionShutdownEvent>): () => void;
on(
event: "session_before_tree",
handler: ExtensionHandler<SessionBeforeTreeEvent, SessionBeforeTreeResult>,
): () => void;
on(event: "session_tree", handler: ExtensionHandler<SessionTreeEvent>): () => void;
on(event: "context", handler: ExtensionHandler<ContextEvent, ContextEventResult>): () => void;
on(
event: "before_provider_request",
handler: ExtensionHandler<BeforeProviderRequestEvent, BeforeProviderRequestEventResult>,
): void;
on(event: "before_provider_headers", handler: ExtensionHandler<BeforeProviderHeadersEvent>): void;
on(event: "after_provider_response", handler: ExtensionHandler<AfterProviderResponseEvent>): void;
on(event: "before_agent_start", handler: ExtensionHandler<BeforeAgentStartEvent, BeforeAgentStartEventResult>): void;
on(event: "agent_start", handler: ExtensionHandler<AgentStartEvent>): void;
on(event: "agent_end", handler: ExtensionHandler<AgentEndEvent>): void;
on(event: "agent_settled", handler: ExtensionHandler<AgentSettledEvent>): void;
on(event: "ui_prompt_start", handler: ExtensionHandler<UIPromptStartEvent>): void;
on(event: "ui_prompt_end", handler: ExtensionHandler<UIPromptEndEvent>): void;
on(event: "turn_start", handler: ExtensionHandler<TurnStartEvent>): void;
on(event: "turn_end", handler: ExtensionHandler<TurnEndEvent>): void;
on(event: "message_start", handler: ExtensionHandler<MessageStartEvent>): void;
on(event: "message_update", handler: ExtensionHandler<MessageUpdateEvent>): void;
on(event: "message_end", handler: ExtensionHandler<MessageEndEvent, MessageEndEventResult>): void;
on(event: "tool_execution_start", handler: ExtensionHandler<ToolExecutionStartEvent>): void;
on(event: "tool_execution_update", handler: ExtensionHandler<ToolExecutionUpdateEvent>): void;
on(event: "tool_execution_end", handler: ExtensionHandler<ToolExecutionEndEvent>): void;
on(event: "model_select", handler: ExtensionHandler<ModelSelectEvent>): void;
on(event: "thinking_level_select", handler: ExtensionHandler<ThinkingLevelSelectEvent>): void;
on(event: "tool_call", handler: ExtensionHandler<ToolCallEvent, ToolCallEventResult>): void;
on(event: "tool_result", handler: ExtensionHandler<ToolResultEvent, ToolResultEventResult>): void;
on(event: "user_bash", handler: ExtensionHandler<UserBashEvent, UserBashEventResult>): void;
on(event: "input", handler: ExtensionHandler<InputEvent, InputEventResult>): void;
): () => void;
on(event: "before_provider_headers", handler: ExtensionHandler<BeforeProviderHeadersEvent>): () => void;
on(event: "after_provider_response", handler: ExtensionHandler<AfterProviderResponseEvent>): () => void;
on(
event: "before_agent_start",
handler: ExtensionHandler<BeforeAgentStartEvent, BeforeAgentStartEventResult>,
): () => void;
on(event: "agent_start", handler: ExtensionHandler<AgentStartEvent>): () => void;
on(event: "agent_end", handler: ExtensionHandler<AgentEndEvent>): () => void;
on(event: "agent_settled", handler: ExtensionHandler<AgentSettledEvent>): () => void;
on(event: "ui_prompt_start", handler: ExtensionHandler<UIPromptStartEvent>): () => void;
on(event: "ui_prompt_end", handler: ExtensionHandler<UIPromptEndEvent>): () => void;
on(event: "turn_start", handler: ExtensionHandler<TurnStartEvent>): () => void;
on(event: "turn_end", handler: ExtensionHandler<TurnEndEvent>): () => void;
on(event: "message_start", handler: ExtensionHandler<MessageStartEvent>): () => void;
on(event: "message_update", handler: ExtensionHandler<MessageUpdateEvent>): () => void;
on(event: "message_end", handler: ExtensionHandler<MessageEndEvent, MessageEndEventResult>): () => void;
on(event: "tool_execution_start", handler: ExtensionHandler<ToolExecutionStartEvent>): () => void;
on(event: "tool_execution_update", handler: ExtensionHandler<ToolExecutionUpdateEvent>): () => void;
on(event: "tool_execution_end", handler: ExtensionHandler<ToolExecutionEndEvent>): () => void;
on(event: "model_select", handler: ExtensionHandler<ModelSelectEvent>): () => void;
on(event: "thinking_level_select", handler: ExtensionHandler<ThinkingLevelSelectEvent>): () => void;
on(event: "tool_call", handler: ExtensionHandler<ToolCallEvent, ToolCallEventResult>): () => void;
on(event: "tool_result", handler: ExtensionHandler<ToolResultEvent, ToolResultEventResult>): () => void;
on(event: "user_bash", handler: ExtensionHandler<UserBashEvent, UserBashEventResult>): () => void;
on(event: "input", handler: ExtensionHandler<InputEvent, InputEventResult>): () => void;
// =========================================================================
// Tool Registration

130
packages/coding-agent/test/extensions-runner.test.ts
View file

@ -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 = `