From 824b7a0a2f0df57ac57abade1fc54b9b7ee195cd Mon Sep 17 00:00:00 2001 From: Federico Jaramillo Martinez Date: Wed, 10 Jun 2026 20:32:25 +0200 Subject: [PATCH] fix: bind extensions for web sessions Fixes #19 --- .changeset/bind-session-extensions.md | 5 ++ src/server/sessions/piSessionService.test.ts | 61 +++++++++++++++++++- src/server/sessions/piSessionService.ts | 27 ++++++++- 3 files changed, 90 insertions(+), 3 deletions(-) create mode 100644 .changeset/bind-session-extensions.md diff --git a/.changeset/bind-session-extensions.md b/.changeset/bind-session-extensions.md new file mode 100644 index 0000000..7d57121 --- /dev/null +++ b/.changeset/bind-session-extensions.md @@ -0,0 +1,5 @@ +--- +"@jmfederico/pi-web": patch +--- + +Initialize Pi extensions for web-managed sessions so `session_start` handlers, extension resources, and startup-dependent tools run correctly. diff --git a/src/server/sessions/piSessionService.test.ts b/src/server/sessions/piSessionService.test.ts index d7133e5..3765859 100644 --- a/src/server/sessions/piSessionService.test.ts +++ b/src/server/sessions/piSessionService.test.ts @@ -45,8 +45,9 @@ function sessionRecord(id: string, cwd = "/workspace") { function fakeRuntime(sessionId = "session-1", patch: Partial = {}) { const promptCalls: { text: string; options: unknown }[] = []; + const bindExtensionCalls: unknown[] = []; const listeners: ((event: unknown) => void)[] = []; - const calls = { abort: 0, clearQueue: 0, dispose: 0, prompt: promptCalls }; + const calls = { abort: 0, bindExtensions: bindExtensionCalls, clearQueue: 0, dispose: 0, prompt: promptCalls }; const session: TestSession = { sessionId, sessionFile: `/tmp/${sessionId}.jsonl`, @@ -71,6 +72,10 @@ function fakeRuntime(sessionId = "session-1", patch: Partial = {}) if (index !== -1) listeners.splice(index, 1); }; }, + bindExtensions: (bindings: unknown) => { + calls.bindExtensions.push(bindings); + return Promise.resolve(); + }, getSessionStats: () => ({ sessionId, totalMessages: 0, userMessages: 0, assistantMessages: 0, toolCalls: 0, tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, cost: 0 }), getContextUsage: () => undefined, prompt: (text: string, options: unknown) => { @@ -146,6 +151,7 @@ describe("PiSessionService", () => { const session = await service.start("/workspace"); expect(createCalls).toBe(1); + expect(fake.calls.bindExtensions).toHaveLength(1); expect(session).toMatchObject({ id: "session-1", cwd: "/workspace", messageCount: 0 }); expect(service.activeCount()).toBe(1); expect(hub.globalEvents.some((event) => event.type === "status.update" && event.status.sessionId === "session-1")).toBe(true); @@ -155,6 +161,59 @@ describe("PiSessionService", () => { expect(fake.calls.dispose).toBe(1); }); + it("binds extensions again when the SDK runtime replaces the active session", async () => { + const hub = new CapturingSessionEventHub(); + const fake = fakeRuntime("session-1"); + const replacement = fakeRuntime("session-2"); + let rebindSession: ((session: PiAgentSession) => Promise) | undefined; + fake.runtime.setRebindSession = (callback) => { rebindSession = callback; }; + const service = new PiSessionService(hub, { + createAgentRuntime: runtimeCreator(fake.runtime), + sessionManager: sessionGateway([]), + heartbeatIntervalMs: 60_000, + }); + + await service.start("/workspace"); + Object.defineProperty(fake.runtime, "session", { configurable: true, value: replacement.session }); + await rebindSession?.(replacement.session); + + expect(fake.calls.bindExtensions).toHaveLength(1); + expect(replacement.calls.bindExtensions).toHaveLength(1); + expect(service.activeCount()).toBe(1); + expect(await service.status("session-2")).toMatchObject({ sessionId: "session-2" }); + + await service.dispose(); + }); + + it("publishes extension errors reported while binding session extensions", async () => { + const hub = new CapturingSessionEventHub(); + const fake = fakeRuntime("extension-session", { + bindExtensions: (bindings) => { + bindings.onError?.({ extensionPath: "pi-mcp-adapter", event: "session_start", error: "MCP failed" }); + return Promise.resolve(); + }, + }); + const service = new PiSessionService(hub, { + createAgentRuntime: runtimeCreator(fake.runtime), + sessionManager: sessionGateway([]), + heartbeatIntervalMs: 60_000, + }); + + await service.start("/workspace"); + + expect(hub.sessionEvents).toContainEqual({ + sessionId: "extension-session", + event: { type: "session.error", message: "pi-mcp-adapter: MCP failed" }, + }); + const extensionErrorActivity = hub.globalEvents.find((event) => event.type === "activity.update" && event.activity.sessionId === "extension-session"); + expect(extensionErrorActivity).toMatchObject({ + type: "activity.update", + activity: { sessionId: "extension-session", phase: "error", label: "extension error", detail: "pi-mcp-adapter: MCP failed" }, + }); + + await service.dispose(); + }); + it("clears stale active activity once a previously active session becomes idle", async () => { vi.useFakeTimers(); let service: PiSessionService | undefined; diff --git a/src/server/sessions/piSessionService.ts b/src/server/sessions/piSessionService.ts index 0f23c6a..e953d79 100644 --- a/src/server/sessions/piSessionService.ts +++ b/src/server/sessions/piSessionService.ts @@ -89,6 +89,17 @@ export interface PiSessionManagerGateway { open(path: string): PiSessionManager; } +interface PiExtensionError { + extensionPath: string; + event: string; + error: string; + stack?: string; +} + +interface PiExtensionBindings { + onError?: (error: PiExtensionError) => void; +} + export interface PiAgentSession { modelRegistry: ModelRegistryInstance; sessionManager: PiSessionManager; @@ -107,6 +118,7 @@ export interface PiAgentSession { promptTemplates: readonly { name: string; description?: string }[]; resourceLoader: { getSkills(): { skills: readonly { name: string; description?: string }[] } }; subscribe(listener: (event: unknown) => void): () => void; + bindExtensions(bindings: PiExtensionBindings): Promise; compact(instructions?: string): Promise<{ summary: string; tokensBefore: number }>; getUserMessagesForForking(): readonly { entryId: string; text: string }[]; getSessionStats(): { sessionId: string; totalMessages: number; userMessages: number; assistantMessages: number; toolCalls: number; tokens: ClientSessionStatus["tokens"]; cost: number }; @@ -642,17 +654,28 @@ export class PiSessionService { private async create(sessionManager: PiSessionManager, cwd: string): Promise> { const runtime = await this.createAgentRuntime(this.createRuntime, { cwd, agentDir: this.agentDir, sessionManager }); + await this.bindSessionExtensions(runtime.session); const active: ActiveSession = { runtime, unsubscribe: noop }; this.bindRuntime(active); - runtime.setRebindSession(() => { + runtime.setRebindSession(async (session) => { + await this.bindSessionExtensions(session); this.bindRuntime(active); - return Promise.resolve(); }); this.active.set(runtime.session.sessionId, active); this.publishStatus(runtime.session); return active; } + private async bindSessionExtensions(session: PiAgentSession): Promise { + await session.bindExtensions({ + onError: (error) => { + const message = `${error.extensionPath}: ${error.error}`; + this.publishActivity(session, "extension error", "error", message); + this.events.publish(session.sessionId, { type: "session.error", message }); + }, + }); + } + private bindRuntime(active: ActiveSession): void { active.unsubscribe(); const { session } = active.runtime;